Coverage for agentos/memory/persistence.py: 19%
190 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 10:19 +0800
1"""
2AgentOS v1.14.9 — Memory Persistence Manager.
4Unified save/load for all 12 memory subsystems, bridging the gap between
5the existing in-memory pyramid and crash-safe disk persistence.
7All writes are atomic (write to temp file, then rename). JSON format with
8gzip compression for production efficiency; plain JSON for debug.
10Usage:
11 mgr = MemoryPersistenceManager(base_dir="~/.agentos/memory")
13 # Save everything
14 await mgr.save_all(
15 pyramid=pyramid,
16 working=working,
17 conversation=conv,
18 long_term=lterm,
19 reflection_engine=reflection,
20 consolidation_pipeline=pipeline,
21 retriever_index=retriever_index,
22 )
24 # Restore everything
25 state = await mgr.load_all()
26"""
28from __future__ import annotations
30import gzip
31import json
32import os
33import tempfile
34import time
35from dataclasses import dataclass, field
36from pathlib import Path
37from typing import Any
39# ── Snapshot Data Models ──────────────────────
42@dataclass
43class MemorySnapshot:
44 """Complete state of all memory subsystems at a point in time."""
46 version: str = "1.14.9"
47 created_at: float = field(default_factory=time.time)
48 # Per-subsystem state dicts (optional — only non-empty ones are saved)
49 pyramid_state: dict[str, Any] = field(default_factory=dict)
50 working_state: dict[str, Any] = field(default_factory=dict)
51 conversation_state: dict[str, Any] = field(default_factory=dict)
52 long_term_state: dict[str, Any] = field(default_factory=dict)
53 reflection_state: dict[str, Any] = field(default_factory=dict)
54 consolidation_state: dict[str, Any] = field(default_factory=dict)
55 retriever_index_state: dict[str, Any] = field(default_factory=dict)
57 def to_dict(self) -> dict[str, Any]:
58 result: dict[str, Any] = {
59 "version": self.version,
60 "created_at": self.created_at,
61 }
62 for field_name in [
63 "pyramid_state",
64 "working_state",
65 "conversation_state",
66 "long_term_state",
67 "reflection_state",
68 "consolidation_state",
69 "retriever_index_state",
70 ]:
71 val = getattr(self, field_name)
72 if val:
73 result[field_name] = val
74 return result
76 @classmethod
77 def from_dict(cls, d: dict[str, Any]) -> MemorySnapshot:
78 return cls(
79 version=d.get("version", "1.14.9"),
80 created_at=d.get("created_at", time.time()),
81 pyramid_state=d.get("pyramid_state", {}),
82 working_state=d.get("working_state", {}),
83 conversation_state=d.get("conversation_state", {}),
84 long_term_state=d.get("long_term_state", {}),
85 reflection_state=d.get("reflection_state", {}),
86 consolidation_state=d.get("consolidation_state", {}),
87 retriever_index_state=d.get("retriever_index_state", {}),
88 )
91# ── Persistence Manager ──────────────────────
94class MemoryPersistenceManager:
95 """Centralized save/load manager for all memory subsystems.
97 Writes snapshots as compressed JSON files under base_dir.
98 Supports atomic writes (temp file + rename) and optional gzip compression.
99 """
101 def __init__(
102 self,
103 base_dir: str = "",
104 compress: bool = True,
105 ):
106 base = Path(base_dir) if base_dir else Path.home() / ".agentos" / "memory"
107 base.mkdir(parents=True, exist_ok=True)
108 self.base_dir: Path = base
109 self.compress = compress
110 self._snapshot_path: Path = base / ("snapshot.json.gz" if compress else "snapshot.json")
111 self._max_backups: int = 3
113 # ── Save ────────────────────────────────
115 async def save_all(
116 self,
117 pyramid: Any = None,
118 working: Any = None,
119 conversation: Any = None,
120 long_term: Any = None,
121 reflection_engine: Any = None,
122 consolidation_pipeline: Any = None,
123 retriever_index: dict[str, Any] | None = None,
124 ) -> str:
125 """Save all memory subsystems to a single snapshot file.
127 Each subsystem provides a get_state() / dump_state() method;
128 we probe for supported interfaces and extract what we can.
130 Returns the snapshot file path.
131 """
132 snapshot = MemorySnapshot()
134 if pyramid is not None:
135 try:
136 snapshot.pyramid_state = pyramid.get_state()
137 except AttributeError:
138 pass
140 if working is not None:
141 try:
142 snapshot.working_state = working.get_state()
143 except AttributeError:
144 pass
146 if conversation is not None:
147 try:
148 snapshot.conversation_state = conversation.get_state()
149 except AttributeError:
150 pass
152 if long_term is not None:
153 try:
154 snapshot.long_term_state = long_term.get_state()
155 except AttributeError:
156 pass
158 if reflection_engine is not None:
159 try:
160 snapshot.reflection_state = reflection_engine.get_state()
161 except AttributeError:
162 pass
164 if consolidation_pipeline is not None:
165 try:
166 snapshot.consolidation_state = consolidation_pipeline.get_state()
167 except AttributeError:
168 pass
170 if retriever_index is not None:
171 snapshot.retriever_index_state = retriever_index
173 return self._atomic_write(snapshot)
175 def save_sync(
176 self,
177 pyramid: Any = None,
178 working: Any = None,
179 conversation: Any = None,
180 long_term: Any = None,
181 reflection_engine: Any = None,
182 consolidation_pipeline: Any = None,
183 retriever_index: dict[str, Any] | None = None,
184 ) -> str:
185 """Synchronous save — for use in signal handlers / atexit hooks."""
186 snapshot = MemorySnapshot()
188 for obj, attr in [
189 (pyramid, "pyramid_state"),
190 (working, "working_state"),
191 (conversation, "conversation_state"),
192 (long_term, "long_term_state"),
193 (reflection_engine, "reflection_state"),
194 (consolidation_pipeline, "consolidation_state"),
195 ]:
196 if obj is not None:
197 try:
198 setattr(snapshot, attr, obj.get_state())
199 except AttributeError:
200 pass
202 if retriever_index is not None:
203 snapshot.retriever_index_state = retriever_index
205 return self._atomic_write(snapshot)
207 # ── Load ────────────────────────────────
209 async def load_all(self) -> MemorySnapshot:
210 """Load the latest memory snapshot from disk.
212 Returns a MemorySnapshot; empty fields mean no saved state for that subsystem.
213 """
214 if not self._snapshot_path.exists():
215 return MemorySnapshot()
217 data = self._read_snapshot_file()
218 if data is None:
219 return MemorySnapshot()
221 return MemorySnapshot.from_dict(data)
223 def load_sync(self) -> MemorySnapshot:
224 """Synchronous load."""
225 if not self._snapshot_path.exists():
226 return MemorySnapshot()
228 data = self._read_snapshot_file()
229 if data is None:
230 return MemorySnapshot()
232 return MemorySnapshot.from_dict(data)
234 async def restore_all(
235 self,
236 pyramid: Any = None,
237 working: Any = None,
238 conversation: Any = None,
239 long_term: Any = None,
240 reflection_engine: Any = None,
241 consolidation_pipeline: Any = None,
242 retriever_index_target: dict[str, Any] | None = None,
243 ) -> int:
244 """Load snapshot from disk and restore into live objects.
246 Each target object must have a restore_state(state_dict) method.
247 Returns count of subsystems restored.
248 """
249 snapshot = await self.load_all()
250 restored = 0
252 for obj, state_attr in [
253 (pyramid, "pyramid_state"),
254 (working, "working_state"),
255 (conversation, "conversation_state"),
256 (long_term, "long_term_state"),
257 (reflection_engine, "reflection_state"),
258 (consolidation_pipeline, "consolidation_state"),
259 ]:
260 state = getattr(snapshot, state_attr, {})
261 if obj is not None and state:
262 try:
263 obj.restore_state(state)
264 restored += 1
265 except AttributeError:
266 pass
268 if retriever_index_target is not None and snapshot.retriever_index_state:
269 retriever_index_target.clear()
270 retriever_index_target.update(snapshot.retriever_index_state)
271 restored += 1
273 return restored
275 # ── Atomic write ────────────────────────
277 def _atomic_write(self, snapshot: MemorySnapshot) -> str:
278 """Write snapshot atomically: temp file → rename."""
279 data = snapshot.to_dict()
280 json_bytes = json.dumps(data, ensure_ascii=False, indent=2, default=str).encode("utf-8")
282 if self.compress:
283 json_bytes = gzip.compress(json_bytes, compresslevel=6)
285 # Write to temp file, then rename
286 fd, tmp_path = tempfile.mkstemp(
287 dir=str(self.base_dir),
288 prefix=".snapshot-tmp-",
289 suffix=".json.gz" if self.compress else ".json",
290 )
291 try:
292 with os.fdopen(fd, "wb") as f:
293 f.write(json_bytes)
295 # Rotate old backups
296 self._rotate_backups()
298 os.replace(tmp_path, str(self._snapshot_path))
299 except Exception:
300 try:
301 os.unlink(tmp_path)
302 except OSError:
303 pass
304 raise
306 return str(self._snapshot_path)
308 # ── Read snapshot ────────────────────────
310 def _read_snapshot_file(self) -> dict[str, Any] | None:
311 """Read and parse snapshot file. Returns None on failure."""
312 try:
313 with open(self._snapshot_path, "rb") as f:
314 raw = f.read()
316 if self.compress:
317 raw = gzip.decompress(raw)
319 return json.loads(raw.decode("utf-8"))
320 except (OSError, json.JSONDecodeError, gzip.BadGzipFile):
321 return None
323 # ── Backup rotation ──────────────────────
325 def _rotate_backups(self) -> None:
326 """Rotate old snapshot backups, keeping self._max_backups."""
327 for i in range(self._max_backups - 1, 0, -1):
328 old_path = self.base_dir / f"snapshot.{i}.json.gz"
329 new_path = self.base_dir / f"snapshot.{i + 1}.json.gz"
330 if old_path.exists():
331 try:
332 os.replace(str(old_path), str(new_path))
333 except OSError:
334 pass
336 # Rotate current into .1
337 if self._snapshot_path.exists():
338 backup_path = self.base_dir / "snapshot.1.json.gz"
339 try:
340 os.replace(str(self._snapshot_path), str(backup_path))
341 except OSError:
342 pass
344 # ── Query ────────────────────────────────
346 def snapshot_info(self) -> dict[str, Any]:
347 """Return metadata about the current snapshot."""
348 if not self._snapshot_path.exists():
349 return {"exists": False}
351 try:
352 stat = self._snapshot_path.stat()
353 snapshot = self.load_sync()
355 subsystems_saved = sum(
356 1
357 for v in [
358 snapshot.pyramid_state,
359 snapshot.working_state,
360 snapshot.conversation_state,
361 snapshot.long_term_state,
362 snapshot.reflection_state,
363 snapshot.consolidation_state,
364 snapshot.retriever_index_state,
365 ]
366 if v
367 )
369 return {
370 "exists": True,
371 "path": str(self._snapshot_path),
372 "size_bytes": stat.st_size,
373 "created_at": snapshot.created_at,
374 "version": snapshot.version,
375 "subsystems_saved": subsystems_saved,
376 "compressed": self.compress,
377 }
378 except Exception:
379 return {"exists": True, "error": "unreadable"}
381 def delete_snapshot(self) -> bool:
382 """Delete the current snapshot file(s)."""
383 deleted = False
384 for path in self.base_dir.glob("snapshot*.json.gz"):
385 try:
386 path.unlink()
387 deleted = True
388 except OSError:
389 pass
390 for path in self.base_dir.glob("snapshot*.json"):
391 try:
392 path.unlink()
393 deleted = True
394 except OSError:
395 pass
396 return deleted
398 def list_backups(self) -> list[dict[str, Any]]:
399 """List all available snapshot backups."""
400 results = []
401 for path in sorted(self.base_dir.glob("snapshot*.json*")):
402 try:
403 stat = path.stat()
404 results.append(
405 {
406 "name": path.name,
407 "size_bytes": stat.st_size,
408 "mtime": stat.st_mtime,
409 }
410 )
411 except OSError:
412 continue
413 return results