Coverage for agentos/memory/session.py: 32%

1"""

2Persistent Thread Context (PTC) — Session Manager.

4OpenClaw-style long-running session management:

5 - Heartbeat: periodic ping to keep sessions alive

6 - Auto-suspend: idle sessions that exceed TTL

7 - State recovery: resume a session exactly where it left off

8 - Cross-session memory: carry context across disconnected sessions

9 - Event hooks: on_suspend, on_resume, on_expire

11Design:

12 SessionManager

13 └─ Session (per-thread lifecycle)

14 ├─ heartbeat() — keep alive

15 ├─ suspend() — save state, pause

16 ├─ resume() — restore state, continue

17 └─ expire() — cleanup after TTL

18"""

20from __future__ import annotations

22import asyncio

23import json

24import time

25import uuid

26from collections.abc import Callable

27from dataclasses import dataclass, field

28from enum import StrEnum

29from pathlib import Path

30from typing import Any

32import aiosqlite

34# ── Session Models ──

37class SessionStatus(StrEnum):

38 ACTIVE = "active" # Currently running

39 IDLE = "idle" # Alive but no recent activity

40 SUSPENDED = "suspended" # Paused, state saved

41 EXPIRED = "expired" # Timed out, cleaned up

42 ERROR = "error" # Crashed but state saved

45@dataclass

46class SessionState:

47 """Serializable state snapshot for a session."""

49 conversation_history: list[dict] = field(default_factory=list)

50 working_memory: dict[str, Any] = field(default_factory=dict)

51 agent_context: dict[str, Any] = field(default_factory=dict)

52 tool_state: dict[str, Any] = field(default_factory=dict)

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

55 def to_json(self) -> str:

56 return json.dumps(

57 {

58 "conversation_history": self.conversation_history,

59 "working_memory": self.working_memory,

60 "agent_context": self.agent_context,

61 "tool_state": self.tool_state,

62 "metadata": self.metadata,

63 },

64 ensure_ascii=False,

65 default=str,

66 )

68 @classmethod

69 def from_json(cls, data: str) -> SessionState:

70 d = json.loads(data)

71 return cls(

72 conversation_history=d.get("conversation_history", []),

73 working_memory=d.get("working_memory", {}),

74 agent_context=d.get("agent_context", {}),

75 tool_state=d.get("tool_state", {}),

76 metadata=d.get("metadata", {}),

77 )

80@dataclass

81class Session:

82 """A single PTC session (one conversational thread)."""

84 id: str = field(default_factory=lambda: uuid.uuid4().hex[:16])

85 name: str = ""

86 user_id: str = "default"

87 status: SessionStatus = SessionStatus.ACTIVE

89 created_at: float = field(default_factory=time.time)

90 last_heartbeat: float = field(default_factory=time.time)

91 last_activity: float = field(default_factory=time.time)

93 state: SessionState = field(default_factory=SessionState)

95 # Config

96 heartbeat_interval: float = 30.0 # seconds between heartbeats

97 idle_timeout: float = 300.0 # idle → suspend (5 min)

98 absolute_ttl: float = 86400.0 # max lifetime (24h)

99 max_history_turns: int = 1000

100

101 # Internal

102 _heartbeat_task: asyncio.Task | None = None

103

104 @property

105 def age_seconds(self) -> float:

106 return time.time() - self.created_at

107

108 @property

109 def idle_seconds(self) -> float:

110 return time.time() - self.last_activity

111

112 @property

113 def is_expired(self) -> bool:

114 return self.age_seconds > self.absolute_ttl

115

116 def to_dict(self) -> dict:

117 return {

118 "id": self.id,

119 "name": self.name,

120 "user_id": self.user_id,

121 "status": self.status.value,

122 "created_at": self.created_at,

123 "last_heartbeat": self.last_heartbeat,

124 "last_activity": self.last_activity,

125 "heartbeat_interval": self.heartbeat_interval,

126 "idle_timeout": self.idle_timeout,

127 "absolute_ttl": self.absolute_ttl,

128 "state": self.state.to_json(),

129 }

130

131

132# ── Session Manager ──

133

134

135class SessionManager:

136 """Manage PTC sessions with heartbeat, suspend/resume, and persistence.

137

138 Usage:

139 manager = SessionManager(db_path="~/.agentos/sessions.db")

140

141 # Create a new session

142 session = await manager.create(name="research-thread")

143

144 # Heartbeat loop (runs in background)

145 await manager.start_heartbeat(session)

146

147 # Suspend on idle

148 await manager.suspend(session.id)

149

150 # Resume later — state restored

151 session = await manager.resume(session.id)

152

153 # Hooks

154 manager.on_suspend(lambda s: print(f"{s.name} suspended"))

155 manager.on_resume(lambda s: print(f"{s.name} resumed"))

156 """

157

158 def __init__(

159 self,

160 db_path: str = "",

161 max_concurrent: int = 100,

162 ):

163 db_path = Path(db_path) if db_path else Path.home() / ".agentos" / "sessions.db"

164 db_path.parent.mkdir(parents=True, exist_ok=True)

165 self._db_path = str(db_path)

166 self._max_concurrent = max_concurrent

167

168 self._sessions: dict[str, Session] = {}

169 self._heartbeat_tasks: dict[str, asyncio.Task] = {}

170 self._hooks: dict[str, list[Callable]] = {

171 "create": [],

172 "suspend": [],

173 "resume": [],

174 "expire": [],

175 "heartbeat_missed": [],

176 }

177

178 # ── Hooks ──

179

180 def on(self, event: str):

181 """Decorator: register a hook for session events."""

182

183 def decorator(fn):

184 self._hooks.setdefault(event, []).append(fn)

185 return fn

186

187 return decorator

188

189 def on_create(self, fn: Callable[[Session], Any]):

190 self._hooks["create"].append(fn)

191

192 def on_suspend(self, fn: Callable[[Session], Any]):

193 self._hooks["suspend"].append(fn)

194

195 def on_resume(self, fn: Callable[[Session], Any]):

196 self._hooks["resume"].append(fn)

197

198 def on_expire(self, fn: Callable[[Session], Any]):

199 self._hooks["expire"].append(fn)

200

201 async def _fire(self, event: str, session: Session):

202 for hook in self._hooks.get(event, []):

203 try:

204 result = hook(session)

205 if asyncio.iscoroutine(result):

206 await result

207 except Exception:

208 pass

209

210 # ── Session Lifecycle ──

211

212 async def create(

213 self,

214 name: str = "",

215 user_id: str = "default",

216 heartbeat_interval: float = 30.0,

217 idle_timeout: float = 300.0,

218 absolute_ttl: float = 86400.0,

219 ) -> Session:

220 """Create a new PTC session."""

221 if len(self._sessions) >= self._max_concurrent:

222 oldest = min(self._sessions.values(), key=lambda s: s.last_activity)

223 await self.expire(oldest.id)

224

225 session = Session(

226 name=name or f"session-{uuid.uuid4().hex[:6]}",

227 user_id=user_id,

228 heartbeat_interval=heartbeat_interval,

229 idle_timeout=idle_timeout,

230 absolute_ttl=absolute_ttl,

231 )

232

233 self._sessions[session.id] = session

234 await self._persist(session)

235 await self._fire("create", session)

236

237 return session

238

239 async def suspend(self, session_id: str) -> bool:

240 """Suspend a session — save state, stop heartbeat."""

241 session = self._sessions.get(session_id)

242 if not session:

243 return False

244

245 session.status = SessionStatus.SUSPENDED

246

247 # Stop heartbeat

248 if session_id in self._heartbeat_tasks:

249 self._heartbeat_tasks[session_id].cancel()

250 del self._heartbeat_tasks[session_id]

251

252 await self._persist(session)

253 await self._fire("suspend", session)

254

255 return True

256

257 async def resume(self, session_id: str) -> Session | None:

258 """Resume a suspended session — restore state, restart heartbeat."""

259 session = self._sessions.get(session_id)

260

261 # Try loading from DB if not in memory

262 if not session:

263 session = await self._load_from_db(session_id)

264 if not session:

265 return None

266

267 if session.status == SessionStatus.EXPIRED:

268 return None

269

270 session.status = SessionStatus.ACTIVE

271 session.last_activity = time.time()

272 session.last_heartbeat = time.time()

273

274 self._sessions[session.id] = session

275 await self._fire("resume", session)

276

277 return session

278

279 async def expire(self, session_id: str) -> bool:

280 """Permanently expire a session — cleanup."""

281 session = self._sessions.pop(session_id, None)

282 if not session:

283 return False

284

285 session.status = SessionStatus.EXPIRED

286

287 if session_id in self._heartbeat_tasks:

288 self._heartbeat_tasks[session_id].cancel()

289 del self._heartbeat_tasks[session_id]

290

291 await self._persist(session)

292 await self._fire("expire", session)

293

294 return True

295

296 async def destroy(self, session_id: str) -> bool:

297 """Hard delete a session from memory and DB."""

298 await self.expire(session_id)

299 async with aiosqlite.connect(self._db_path) as db:

300 await db.execute("DELETE FROM sessions WHERE id = ?", (session_id,))

301 await db.commit()

302 return True

303

304 # ── Heartbeat ──

305

306 async def start_heartbeat(self, session: Session) -> None:

307 """Start background heartbeat for a session."""

308 if session.id in self._heartbeat_tasks:

309 return

310

311 async def _loop():

312 while True:

313 await asyncio.sleep(session.heartbeat_interval)

314

315 if session.id not in self._sessions:

316 return

317

318 session.last_heartbeat = time.time()

319

320 # Check idle timeout → suspend

321 if session.idle_seconds > session.idle_timeout:

322 await self.suspend(session.id)

323 return

324

325 # Check absolute TTL → expire

326 if session.is_expired:

327 await self.expire(session.id)

328 return

329

330 # Re-persist state snapshot

331 await self._persist(session)

332

333 self._heartbeat_tasks[session.id] = asyncio.create_task(_loop())

334

335 async def heartbeat(self, session_id: str) -> bool:

336 """Manual heartbeat ping. Returns False if session not found."""

337 session = self._sessions.get(session_id)

338 if not session:

339 return False

340

341 session.last_heartbeat = time.time()

342 session.last_activity = time.time()

343

344 # Re-activate if suspended

345 if session.status == SessionStatus.SUSPENDED:

346 await self.resume(session_id)

347

348 return True

349

350 # ── State Management ──

351

352 async def save_state(self, session_id: str, state: SessionState) -> bool:

353 """Save explicit state snapshot for a session."""

354 session = self._sessions.get(session_id)

355 if not session:

356 return False

357

358 session.state = state

359 session.last_activity = time.time()

360 await self._persist(session)

361 return True

362

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

364 """Get the latest state snapshot for a session."""

365 session = self._sessions.get(session_id)

366 if session:

367 return session.state

368

369 session = await self._load_from_db(session_id)

370 return session.state if session else None

371

372 async def add_context(self, session_id: str, key: str, value: Any) -> bool:

373 """Add a key-value to the session's working memory."""

374 session = self._sessions.get(session_id)

375 if not session:

376 return False

377 session.state.working_memory[key] = value

378 session.last_activity = time.time()

379 return True

380

381 # ── Query ──

382

383 def get(self, session_id: str) -> Session | None:

384 """Get an active session by ID."""

385 return self._sessions.get(session_id)

386

387 def list_active(self, user_id: str = "") -> list[Session]:

388 """List all active/idle sessions, optionally filtered by user."""

389 sessions = [

390 s

391 for s in self._sessions.values()

392 if s.status in (SessionStatus.ACTIVE, SessionStatus.IDLE)

393 ]

394 if user_id:

395 sessions = [s for s in sessions if s.user_id == user_id]

396 return sorted(sessions, key=lambda s: s.last_activity, reverse=True)

397

398 def list_suspended(self, user_id: str = "") -> list[Session]:

399 """List suspended sessions."""

400 sessions = [s for s in self._sessions.values() if s.status == SessionStatus.SUSPENDED]

401 if user_id:

402 sessions = [s for s in sessions if s.user_id == user_id]

403 return sorted(sessions, key=lambda s: s.last_activity, reverse=True)

404

405 async def count(self) -> int:

406 """Total sessions in memory."""

407 return len(self._sessions)

408

409 # ── Monitor ──

410

411 async def monitor(self) -> dict[str, Any]:

412 """Get a monitoring snapshot of all sessions."""

413 active = 0

414 idle = 0

415 suspended = 0

416

417 for s in self._sessions.values():

418 if s.status == SessionStatus.ACTIVE:

419 active += 1

420 elif s.status == SessionStatus.IDLE:

421 idle += 1

422 elif s.status == SessionStatus.SUSPENDED:

423 suspended += 1

424

425 return {

426 "total": len(self._sessions),

427 "active": active,

428 "idle": idle,

429 "suspended": suspended,

430 "heartbeat_tasks": len(self._heartbeat_tasks),

431 }

432

433 # ── Persistence ──

434

435 async def _persist(self, session: Session) -> None:

436 """Save session to SQLite."""

437 try:

438 async with aiosqlite.connect(self._db_path) as db:

439 await db.execute("""

440 CREATE TABLE IF NOT EXISTS sessions (

441 id TEXT PRIMARY KEY,

442 name TEXT,

443 user_id TEXT,

444 status TEXT,

445 created_at REAL,

446 last_heartbeat REAL,

447 last_activity REAL,

448 heartbeat_interval REAL,

449 idle_timeout REAL,

450 absolute_ttl REAL,

451 state TEXT

452 )

453 """)

454 await db.execute(

455 """

456 INSERT OR REPLACE INTO sessions

457 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)

458 """,

459 (

460 session.id,

461 session.name,

462 session.user_id,

463 session.status.value,

464 session.created_at,

465 session.last_heartbeat,

466 session.last_activity,

467 session.heartbeat_interval,

468 session.idle_timeout,

469 session.absolute_ttl,

470 session.state.to_json(),

471 ),

472 )

473 await db.commit()

474 except Exception:

475 pass

476

477 async def _load_from_db(self, session_id: str) -> Session | None:

478 """Load a session from SQLite."""

479 try:

480 async with aiosqlite.connect(self._db_path) as db:

481 await db.execute("""

482 CREATE TABLE IF NOT EXISTS sessions (

483 id TEXT PRIMARY KEY, name TEXT, user_id TEXT,

484 status TEXT, created_at REAL, last_heartbeat REAL,

485 last_activity REAL, heartbeat_interval REAL,

486 idle_timeout REAL, absolute_ttl REAL, state TEXT

487 )

488 """)

489 cursor = await db.execute("SELECT * FROM sessions WHERE id = ?", (session_id,))

490 row = await cursor.fetchone()

491 if not row:

492 return None

493

494 session = Session(

495 id=row[0],

496 name=row[1],

497 user_id=row[2],

498 status=SessionStatus(row[3]),

499 created_at=row[4],

500 last_heartbeat=row[5],

501 last_activity=row[6],

502 heartbeat_interval=row[7],

503 idle_timeout=row[8],

504 absolute_ttl=row[9],

505 state=SessionState.from_json(row[10]),

506 )

507 self._sessions[session.id] = session

508 return session

509 except Exception:

510 return None