Coverage for merco/memory/session_store.py: 84%
167 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 14:04 +0800
1"""SQLite 会话持久化存储"""
3import json
4import logging
5import os
6import shutil
7import sqlite3
8import time
9import uuid
10from datetime import datetime
11from pathlib import Path
13logger = logging.getLogger("merco.session_store")
16class SessionWriteError(Exception):
17 """Session 写入失败（重试 N 次后仍失败）"""
18 pass
21class SessionStore:
22 """SQLite 会话存储 — 两张表，自动建表，增量写消息"""
24 def __init__(self, db_path: str):
25 self.db_path = os.path.expanduser(db_path)
26 # 顺序：先检查/恢复，再建表
27 # （如果 DB 完全损坏，_ensure_db() 会在 _conn() 阶段抛 DatabaseError）
28 self.startup_check()
29 self._ensure_db()
31 # ── 表初始化 ──────────────────────────────────────────
33 def _ensure_db(self):
34 os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
35 with self._conn() as conn:
36 conn.executescript("""
37 CREATE TABLE IF NOT EXISTS sessions (
38 id TEXT PRIMARY KEY,
39 title TEXT DEFAULT '',
40 created_at TEXT NOT NULL,
41 updated_at TEXT NOT NULL,
42 message_count INTEGER DEFAULT 0,
43 parent_id TEXT,
44 metadata TEXT DEFAULT '{}'
45 );
47 CREATE TABLE IF NOT EXISTS messages (
48 id INTEGER PRIMARY KEY AUTOINCREMENT,
49 session_id TEXT NOT NULL,
50 role TEXT NOT NULL,
51 content TEXT DEFAULT '',
52 tool_call_id TEXT DEFAULT '',
53 tool_calls TEXT DEFAULT '[]',
54 reasoning TEXT DEFAULT '',
55 timestamp TEXT NOT NULL,
56 FOREIGN KEY (session_id) REFERENCES sessions(id)
57 );
59 CREATE INDEX IF NOT EXISTS idx_msg_session
60 ON messages(session_id, id);
61 """)
62 # 兼容已有数据库：加 metadata 列
63 try:
64 conn.execute("ALTER TABLE sessions ADD COLUMN metadata TEXT DEFAULT '{}'")
65 except Exception:
66 pass # 列已存在
68 def _conn(self) -> sqlite3.Connection:
69 conn = sqlite3.connect(self.db_path)
70 conn.execute("PRAGMA journal_mode=WAL")
71 conn.execute("PRAGMA foreign_keys=ON")
72 conn.row_factory = sqlite3.Row
73 return conn
75 # ── 备份与恢复 ──────────────────────────────────────
77 @property
78 def _backup_path(self) -> str:
79 """备份文件路径"""
80 return self.db_path + ".backup"
82 def backup(self) -> bool:
83 """创建数据库备份，返回 True=成功
85 Note: 先执行 WAL checkpoint 确保所有 pending 事务都刷入主数据库文件
86 """
87 try:
88 # 强制 WAL 刷盘，避免备份丢失未提交的事务
89 with self._conn() as conn:
90 conn.execute("PRAGMA wal_checkpoint(FULL)")
91 shutil.copy2(self.db_path, self._backup_path)
92 logger.debug(f"Session 数据库已备份到 {self._backup_path}")
93 return True
94 except Exception as e:
95 logger.warning(f"备份 Session 数据库失败: {e}")
96 return False
98 def restore_from_backup(self) -> bool:
99 """从备份恢复数据库，返回 True=成功"""
100 if not os.path.exists(self._backup_path):
101 logger.warning("无备份文件，无法恢复")
102 return False
103 try:
104 shutil.copy2(self._backup_path, self.db_path)
105 logger.info("从备份恢复 Session 数据库成功")
106 return True
107 except Exception as e:
108 logger.error(f"从备份恢复失败: {e}")
109 return False
111 def delete_backup(self) -> bool:
112 """删除备份文件，返回 True=成功"""
113 if not os.path.exists(self._backup_path):
114 return True # 文件不存在，视为成功
115 try:
116 os.remove(self._backup_path)
117 logger.debug("备份文件已删除")
118 return True
119 except Exception as e:
120 logger.warning(f"删除备份文件失败: {e}")
121 return False
123 # ── 完整性检查 ──────────────────────────────────────
125 def check_integrity(self) -> bool:
126 """检查 SQLite 完整性，返回 True=正常"""
127 try:
128 with self._conn() as conn:
129 result = conn.execute("PRAGMA integrity_check").fetchone()
130 is_ok = result[0] == "ok"
131 if not is_ok:
132 logger.warning(f"Session 数据库完整性检查失败: {result[0]}")
133 return is_ok
134 except Exception as e:
135 logger.warning(f"完整性检查异常: {e}")
136 return False
138 def startup_check(self):
139 """启动时检查，必要时恢复"""
140 if not os.path.exists(self.db_path):
141 return # 数据库不存在，无需检查
143 if not self.check_integrity():
144 logger.warning("Session 数据库损坏，尝试从备份恢复...")
145 if not self.restore_from_backup():
146 logger.error("恢复失败，Session 数据可能丢失")
147 # 无备份可用：删除损坏文件，让 _ensure_db() 重建
148 try:
149 os.remove(self.db_path)
150 logger.warning("已删除损坏的 Session 数据库")
151 except Exception as e:
152 logger.error(f"删除损坏文件失败: {e}")
154 # ── Session CRUD ──────────────────────────────────────
156 def create_session(self, session_id: str, title: str = "") -> dict:
157 now = _now()
158 with self._conn() as conn:
159 conn.execute(
160 "INSERT OR IGNORE INTO sessions (id, title, created_at, updated_at) "
161 "VALUES (?, ?, ?, ?)",
162 (session_id, title or "", now, now),
163 )
164 conn.commit()
165 return {"id": session_id, "title": title, "created_at": now}
167 def save_message(self, session_id: str, role: str, content: str = "",
168 tool_call_id: str = "", tool_calls: list | None = None,
169 reasoning: str = "") -> int:
170 """保存消息，支持重试"""
171 now = _now()
172 tc_json = json.dumps(tool_calls or [], ensure_ascii=False)
173 last_error = None
175 for attempt in range(3):
176 try:
177 with self._conn() as conn:
178 cur = conn.execute(
179 "INSERT INTO messages (session_id, role, content, tool_call_id, "
180 "tool_calls, reasoning, timestamp) "
181 "VALUES (?, ?, ?, ?, ?, ?, ?)",
182 (session_id, role, content, tool_call_id, tc_json, reasoning, now),
183 )
184 conn.execute(
185 "UPDATE sessions SET updated_at = ?, "
186 "message_count = message_count + 1 "
187 "WHERE id = ?",
188 (now, session_id),
189 )
190 conn.commit()
191 return cur.lastrowid
192 except sqlite3.OperationalError as e:
193 last_error = e
194 if attempt < 2:
195 time.sleep(0.1 * (attempt + 1))
196 except Exception as e:
197 # 非瞬时错误，不重试
198 raise SessionWriteError(f"写入失败（非重试错误）: {e}") from e
200 # 3 次全部失败
201 logger.warning(
202 f"⚠️ Session 写入失败（已重试 3 次）\n"
203 f" 可能是磁盘满或权限问题\n"
204 f" 建议：检查 ~/.merco/ 目录"
205 )
206 raise SessionWriteError(f"写入失败: {last_error}")
208 def load_session(self, session_id: str) -> dict | None:
209 with self._conn() as conn:
210 row = conn.execute(
211 "SELECT * FROM sessions WHERE id = ?", (session_id,)
212 ).fetchone()
213 if not row:
214 return None
216 messages = conn.execute(
217 "SELECT * FROM messages WHERE session_id = ? ORDER BY id",
218 (session_id,),
219 ).fetchall()
221 return {
222 "id": row["id"],
223 "title": row["title"],
224 "created_at": row["created_at"],
225 "updated_at": row["updated_at"],
226 "message_count": row["message_count"],
227 "parent_id": row["parent_id"],
228 "metadata": json.loads(row["metadata"] or "{}"),
229 "messages": [
230 {
231 "role": m["role"],
232 "content": m["content"],
233 "tool_call_id": m["tool_call_id"],
234 "tool_calls": json.loads(m["tool_calls"]),
235 "reasoning": m["reasoning"],
236 "timestamp": m["timestamp"],
237 }
238 for m in messages
239 ],
240 }
242 def list_sessions(self, limit: int = 20) -> list[dict]:
243 with self._conn() as conn:
244 rows = conn.execute(
245 "SELECT id, title, created_at, updated_at, message_count "
246 "FROM sessions ORDER BY updated_at DESC LIMIT ?",
247 (limit,),
248 ).fetchall()
249 return [dict(r) for r in rows]
251 def count_messages(self, session_id: str) -> int:
252 with self._conn() as conn:
253 row = conn.execute(
254 "SELECT COUNT(*) as cnt FROM messages WHERE session_id = ?",
255 (session_id,),
256 ).fetchone()
257 return row["cnt"] if row else 0
259 def save_metadata(self, session_id: str, metadata: dict):
260 import json
261 with self._conn() as conn:
262 conn.execute(
263 "UPDATE sessions SET metadata = ? WHERE id = ?",
264 (json.dumps(metadata, ensure_ascii=False), session_id),
265 )
266 conn.commit()
268 def update_title(self, session_id: str, title: str):
269 with self._conn() as conn:
270 conn.execute(
271 "UPDATE sessions SET title = ? WHERE id = ? AND title = ''",
272 (title, session_id),
273 )
274 conn.commit()
276 def set_title(self, session_id: str, title: str) -> bool:
277 """Always set the title, regardless of current value. Returns True if a row was updated."""
278 with self._conn() as conn:
279 cur = conn.execute(
280 "UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?",
281 (title, _now(), session_id),
282 )
283 conn.commit()
284 return cur.rowcount > 0
286 def delete_session(self, session_id: str):
287 with self._conn() as conn:
288 conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
289 conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
290 conn.commit()
292 def clone_session(self, session_id: str) -> str:
293 """Clone a session and all its messages to a new session ID.
295 Returns the new session ID (uuid4 hex string).
296 Raises ValueError if the original session does not exist.
297 """
298 original = self.load_session(session_id)
299 if original is None:
300 raise ValueError(f"Session not found: {session_id}")
302 new_id = uuid.uuid4().hex
303 now = _now()
305 with self._conn() as conn:
306 conn.execute(
307 "INSERT INTO sessions (id, title, created_at, updated_at, "
308 "message_count, parent_id, metadata) "
309 "VALUES (?, ?, ?, ?, ?, ?, ?)",
310 (
311 new_id,
312 original["title"],
313 now,
314 now,
315 0, # updated atomically after message copies
316 session_id,
317 json.dumps(original["metadata"], ensure_ascii=False),
318 ),
319 )
321 for msg in original["messages"]:
322 tc_json = json.dumps(msg.get("tool_calls") or [], ensure_ascii=False)
323 conn.execute(
324 "INSERT INTO messages (session_id, role, content, "
325 "tool_call_id, tool_calls, reasoning, timestamp) "
326 "VALUES (?, ?, ?, ?, ?, ?, ?)",
327 (
328 new_id,
329 msg["role"],
330 msg.get("content", ""),
331 msg.get("tool_call_id", ""),
332 tc_json,
333 msg.get("reasoning", ""),
334 now,
335 ),
336 )
338 # Update message_count atomically
339 conn.execute(
340 "UPDATE sessions SET message_count = ("
341 "SELECT COUNT(*) FROM messages WHERE session_id = ?"
342 ") WHERE id = ?",
343 (new_id, new_id),
344 )
345 conn.commit()
347 return new_id
349 def get_children(self, session_id: str) -> list[dict]:
350 with self._conn() as conn:
351 rows = conn.execute(
352 "SELECT id, title, created_at, message_count "
353 "FROM sessions WHERE parent_id = ? ORDER BY created_at DESC",
354 (session_id,)
355 ).fetchall()
356 return [dict(r) for r in rows]
359def _now() -> str:
360 return datetime.now().isoformat(timespec="seconds")