Compare commits

...
Author SHA1 Message Date
alvinfeiandCursor 35f86dc96d fix(sdk): serialize sqlite connection lifecycle
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-19 12:03:29 +08:00
4 changed files with 355 additions and 195 deletions

No files matched your search

+2
View File
@@ -14,6 +14,8 @@ Hy Memory and MemHub are versioned and released independently.
## Unreleased ## Unreleased
- Serialized each SDK SQLite store's shared connection lifecycle so concurrent
history/cache operations cannot race with `close()` and crash the process.
- Clarified that `add(metadata=...)` carries legacy SDK processing options, - Clarified that `add(metadata=...)` carries legacy SDK processing options,
while persistent caller business context belongs in `meta`. while persistent caller business context belongs in `meta`.
- Added an opt-in `get(..., include_raw=True)` path for reading `L1_RAW` - Added an opt-in `get(..., include_raw=True)` path for reading `L1_RAW`
+136 -125
View File
@@ -7,7 +7,7 @@ Agent Memory V2 - SqliteCache
设计模式跟随 history_store.py: 设计模式跟随 history_store.py:
- 标准库 sqlite3 + _run_in_sqlite_pool() - 标准库 sqlite3 + _run_in_sqlite_pool()
- WAL 模式 + check_same_thread=False - WAL 模式 + check_same_thread=False
- threading.Lock 保护写操作 - 每实例 threading.RLock 保护共享连接的全部操作与关闭
三张表: 三张表:
- memory_operations / pipeline_logs / system_metrics: 审计与观测 - memory_operations / pipeline_logs / system_metrics: 审计与观测
@@ -122,7 +122,10 @@ class SqliteCache(CacheBase):
def __init__(self, config: MemoryConfig): def __init__(self, config: MemoryConfig):
self.config = config self.config = config
self._conn: Optional[sqlite3.Connection] = None self._conn: Optional[sqlite3.Connection] = None
self._lock = threading.Lock() # One connection is shared by the executor workers. Use the same lock
# for reads, writes, initialization and close so close() cannot free
# the connection while another worker is inside sqlite3.
self._lock = threading.RLock()
# 读取 db_path # 读取 db_path
cache_cfg = getattr(config, "cache", None) cache_cfg = getattr(config, "cache", None)
@@ -144,35 +147,37 @@ class SqliteCache(CacheBase):
logger.debug(f"SqliteCache initialized: {self._db_path}") logger.debug(f"SqliteCache initialized: {self._db_path}")
def _init_sync(self) -> None: def _init_sync(self) -> None:
db_dir = os.path.dirname(self._db_path) with self._lock:
if db_dir: db_dir = os.path.dirname(self._db_path)
os.makedirs(db_dir, exist_ok=True) if db_dir:
self._conn = sqlite3.connect(self._db_path, check_same_thread=False) os.makedirs(db_dir, exist_ok=True)
self._conn.row_factory = sqlite3.Row self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
self._conn.execute("PRAGMA journal_mode=WAL;") self._conn.row_factory = sqlite3.Row
self._conn.executescript(_CREATE_TABLES_SQL) self._conn.execute("PRAGMA journal_mode=WAL;")
self._migrate_pipeline_logs_columns() self._conn.executescript(_CREATE_TABLES_SQL)
self._conn.commit() self._migrate_pipeline_logs_columns()
self._conn.commit()
def _migrate_pipeline_logs_columns(self) -> None: def _migrate_pipeline_logs_columns(self) -> None:
"""给旧库的 pipeline_logs_v2 补 namespace/session/event time 列。""" """给旧库的 pipeline_logs_v2 补 namespace/session/event time 列。"""
try: with self._lock:
cols = { try:
row["name"] cols = {
for row in self._conn.execute("PRAGMA table_info(pipeline_logs_v2)") row["name"]
} for row in self._conn.execute("PRAGMA table_info(pipeline_logs_v2)")
_text_cols = ("namespace", "session_id") }
for col in _text_cols: _text_cols = ("namespace", "session_id")
if col not in cols: for col in _text_cols:
if col not in cols:
self._conn.execute(
f"ALTER TABLE pipeline_logs_v2 ADD COLUMN {col} TEXT NOT NULL DEFAULT ''"
)
if "event_at_ms" not in cols:
self._conn.execute( self._conn.execute(
f"ALTER TABLE pipeline_logs_v2 ADD COLUMN {col} TEXT NOT NULL DEFAULT ''" "ALTER TABLE pipeline_logs_v2 ADD COLUMN event_at_ms INTEGER NOT NULL DEFAULT 0"
) )
if "event_at_ms" not in cols: except Exception as e:
self._conn.execute( logger.debug(f"[sqlite] pipeline_logs_v2 column migration skipped: {e}")
"ALTER TABLE pipeline_logs_v2 ADD COLUMN event_at_ms INTEGER NOT NULL DEFAULT 0"
)
except Exception as e:
logger.debug(f"[sqlite] pipeline_logs_v2 column migration skipped: {e}")
async def close(self) -> None: async def close(self) -> None:
if self._conn is not None: if self._conn is not None:
@@ -180,9 +185,10 @@ class SqliteCache(CacheBase):
logger.info("SqliteCache closed") logger.info("SqliteCache closed")
def _close_sync(self) -> None: def _close_sync(self) -> None:
if self._conn: with self._lock:
self._conn.close() if self._conn:
self._conn = None self._conn.close()
self._conn = None
# ================================================================ # ================================================================
# 统计 # 统计
@@ -255,39 +261,40 @@ class SqliteCache(CacheBase):
""" """
try: try:
def _query(): def _query():
conditions = [] with self._lock:
params = [] conditions = []
if request_id: params = []
conditions.append("request_id = ?") if request_id:
params.append(request_id) conditions.append("request_id = ?")
if memory_id: params.append(request_id)
conditions.append("memory_id = ?") if memory_id:
params.append(memory_id) conditions.append("memory_id = ?")
if user_id: params.append(memory_id)
conditions.append("user_id = ?") if user_id:
params.append(user_id) conditions.append("user_id = ?")
params.append(user_id)
if not conditions: if not conditions:
return [] return []
where = " AND ".join(conditions) where = " AND ".join(conditions)
rows = self._conn.execute( rows = self._conn.execute(
f"SELECT * FROM memory_operations_v2 WHERE {where} " f"SELECT * FROM memory_operations_v2 WHERE {where} "
f"ORDER BY id DESC LIMIT ?", f"ORDER BY id DESC LIMIT ?",
params + [limit], params + [limit],
).fetchall() ).fetchall()
result = [] result = []
for r in rows: for r in rows:
row = dict(r) row = dict(r)
# 反序列化 supersedes JSON 字符串 → list # 反序列化 supersedes JSON 字符串 → list
import json as _json import json as _json
raw_sup = row.get("supersedes", "") raw_sup = row.get("supersedes", "")
try: try:
row["supersedes"] = _json.loads(raw_sup) if raw_sup else [] row["supersedes"] = _json.loads(raw_sup) if raw_sup else []
except Exception: except Exception:
row["supersedes"] = [] row["supersedes"] = []
result.append(row) result.append(row)
return result return result
return await _run_in_sqlite_pool(_query) return await _run_in_sqlite_pool(_query)
except Exception as e: except Exception as e:
@@ -369,37 +376,38 @@ class SqliteCache(CacheBase):
""" """
try: try:
def _query(): def _query():
conditions = [] with self._lock:
params = [] conditions = []
if request_id: params = []
conditions.append("request_id = ?") if request_id:
params.append(request_id) conditions.append("request_id = ?")
if user_id: params.append(request_id)
conditions.append("user_id = ?") if user_id:
params.append(user_id) conditions.append("user_id = ?")
if step: params.append(user_id)
conditions.append("step = ?") if step:
params.append(step) conditions.append("step = ?")
params.append(step)
if not conditions: if not conditions:
return [] return []
where = " AND ".join(conditions) where = " AND ".join(conditions)
rows = self._conn.execute( rows = self._conn.execute(
f"SELECT * FROM pipeline_logs_v2 WHERE {where} " f"SELECT * FROM pipeline_logs_v2 WHERE {where} "
f"ORDER BY id ASC LIMIT ?", f"ORDER BY id ASC LIMIT ?",
params + [limit], params + [limit],
).fetchall() ).fetchall()
out = [dict(r) for r in rows] out = [dict(r) for r in rows]
# 派生可读的 event_at(ISO ms)供查询接口直接用;event_at_ms 为 0 时留空 # 派生可读的 event_at(ISO ms)供查询接口直接用;event_at_ms 为 0 时留空
from datetime import datetime as _dt from datetime import datetime as _dt
for _r in out: for _r in out:
_ms = _r.get("event_at_ms") or 0 _ms = _r.get("event_at_ms") or 0
_r["event_at"] = ( _r["event_at"] = (
_dt.fromtimestamp(_ms / 1000.0).isoformat(timespec="milliseconds") _dt.fromtimestamp(_ms / 1000.0).isoformat(timespec="milliseconds")
if _ms else "" if _ms else ""
) )
return out return out
return await _run_in_sqlite_pool(_query) return await _run_in_sqlite_pool(_query)
except Exception as e: except Exception as e:
@@ -414,26 +422,27 @@ class SqliteCache(CacheBase):
"""存储一个分钟桶的增量指标数据(UPSERT: 同一分钟多次写入会合并)""" """存储一个分钟桶的增量指标数据(UPSERT: 同一分钟多次写入会合并)"""
try: try:
def _store(): def _store():
now = datetime.now().isoformat() with self._lock:
# UPSERT: 如果同一分钟已存在,合并数据 now = datetime.now().isoformat()
existing = self._conn.execute( # UPSERT: 如果同一分钟已存在,合并数据
"SELECT data FROM system_metrics WHERE minute_ts = ?", existing = self._conn.execute(
(minute_ts,), "SELECT data FROM system_metrics WHERE minute_ts = ?",
).fetchone() (minute_ts,),
if existing: ).fetchone()
import json as _json if existing:
old_data = _json.loads(existing["data"]) import json as _json
merged = self._merge_metric_buckets(old_data, data) old_data = _json.loads(existing["data"])
self._conn.execute( merged = self._merge_metric_buckets(old_data, data)
"UPDATE system_metrics SET data = ?, created_at = ? WHERE minute_ts = ?", self._conn.execute(
(_json.dumps(merged, default=str), now, minute_ts), "UPDATE system_metrics SET data = ?, created_at = ? WHERE minute_ts = ?",
) (_json.dumps(merged, default=str), now, minute_ts),
else: )
self._conn.execute( else:
"INSERT INTO system_metrics (minute_ts, data, created_at) VALUES (?, ?, ?)", self._conn.execute(
(minute_ts, json.dumps(data, default=str), now), "INSERT INTO system_metrics (minute_ts, data, created_at) VALUES (?, ?, ?)",
) (minute_ts, json.dumps(data, default=str), now),
self._conn.commit() )
self._conn.commit()
await _run_in_sqlite_pool(_store) await _run_in_sqlite_pool(_store)
except Exception as e: except Exception as e:
logger.warning(f"store_metrics_minute failed: {e}") logger.warning(f"store_metrics_minute failed: {e}")
@@ -442,17 +451,18 @@ class SqliteCache(CacheBase):
"""读取指定时间范围内的分钟指标数据""" """读取指定时间范围内的分钟指标数据"""
try: try:
def _load(): def _load():
rows = self._conn.execute( with self._lock:
"SELECT data FROM system_metrics WHERE minute_ts >= ? AND minute_ts <= ? ORDER BY minute_ts", rows = self._conn.execute(
(start_ts, end_ts), "SELECT data FROM system_metrics WHERE minute_ts >= ? AND minute_ts <= ? ORDER BY minute_ts",
).fetchall() (start_ts, end_ts),
results = [] ).fetchall()
for row in rows: results = []
try: for row in rows:
results.append(json.loads(row["data"])) try:
except (json.JSONDecodeError, TypeError): results.append(json.loads(row["data"]))
pass except (json.JSONDecodeError, TypeError):
return results pass
return results
return await _run_in_sqlite_pool(_load) return await _run_in_sqlite_pool(_load)
except Exception as e: except Exception as e:
logger.warning(f"load_metrics_range failed: {e}") logger.warning(f"load_metrics_range failed: {e}")
@@ -462,11 +472,12 @@ class SqliteCache(CacheBase):
"""删除指定时间之前的 metrics 数据""" """删除指定时间之前的 metrics 数据"""
try: try:
def _cleanup(): def _cleanup():
self._conn.execute( with self._lock:
"DELETE FROM system_metrics WHERE minute_ts < ?", self._conn.execute(
(before_ts,), "DELETE FROM system_metrics WHERE minute_ts < ?",
) (before_ts,),
self._conn.commit() )
self._conn.commit()
await _run_in_sqlite_pool(_cleanup) await _run_in_sqlite_pool(_cleanup)
except Exception as e: except Exception as e:
logger.warning(f"cleanup_old_metrics failed: {e}") logger.warning(f"cleanup_old_metrics failed: {e}")
+81 -70
View File
@@ -16,6 +16,7 @@ import json
import logging import logging
import os import os
import sqlite3 import sqlite3
import threading
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Dict, List, Any, Optional from typing import Dict, List, Any, Optional
@@ -89,6 +90,10 @@ class HistoryStore:
""" """
self._db_path = config.history.db_path self._db_path = config.history.db_path
self._conn: Optional[sqlite3.Connection] = None self._conn: Optional[sqlite3.Connection] = None
# sqlite3.Connection.close() must never race with an execute() on the
# same connection. The lock also serializes operations on this shared
# connection; separate HistoryStore instances remain independent.
self._conn_lock = threading.RLock()
async def initialize(self) -> None: async def initialize(self) -> None:
"""创建表和索引(如果不存在)""" """创建表和索引(如果不存在)"""
@@ -97,20 +102,21 @@ class HistoryStore:
def _init_sync(self) -> None: def _init_sync(self) -> None:
"""同步初始化:创建目录、连接数据库、建表建索引""" """同步初始化:创建目录、连接数据库、建表建索引"""
# 确保目录存在 with self._conn_lock:
db_dir = os.path.dirname(self._db_path) # 确保目录存在
if db_dir: db_dir = os.path.dirname(self._db_path)
os.makedirs(db_dir, exist_ok=True) if db_dir:
os.makedirs(db_dir, exist_ok=True)
self._conn = sqlite3.connect(self._db_path, check_same_thread=False) self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
self._conn.row_factory = sqlite3.Row self._conn.row_factory = sqlite3.Row
# 开启 WAL 模式,提高并发读写性能 # 开启 WAL 模式,提高并发读写性能
self._conn.execute("PRAGMA journal_mode=WAL;") self._conn.execute("PRAGMA journal_mode=WAL;")
self._conn.execute(_CREATE_TABLE_SQL) self._conn.execute(_CREATE_TABLE_SQL)
for idx_sql in _CREATE_INDEXES_SQL: for idx_sql in _CREATE_INDEXES_SQL:
self._conn.execute(idx_sql) self._conn.execute(idx_sql)
self._conn.commit() self._conn.commit()
# ================================================================ # ================================================================
# 记录操作 # 记录操作
@@ -218,13 +224,14 @@ class HistoryStore:
return await _run_in_sqlite_pool(self._get_history_sync, memory_id) return await _run_in_sqlite_pool(self._get_history_sync, memory_id)
def _get_history_sync(self, memory_id: str) -> List[Dict[str, Any]]: def _get_history_sync(self, memory_id: str) -> List[Dict[str, Any]]:
if self._conn is None: with self._conn_lock:
return [] if self._conn is None:
cursor = self._conn.execute( return []
"SELECT * FROM memory_history WHERE memory_id = ? ORDER BY created_at ASC", cursor = self._conn.execute(
(memory_id,), "SELECT * FROM memory_history WHERE memory_id = ? ORDER BY created_at ASC",
) (memory_id,),
return [dict(row) for row in cursor.fetchall()] )
return [dict(row) for row in cursor.fetchall()]
async def get_recent( async def get_recent(
self, self,
@@ -244,14 +251,15 @@ class HistoryStore:
return await _run_in_sqlite_pool(self._get_recent_sync, isolation_key, limit) return await _run_in_sqlite_pool(self._get_recent_sync, isolation_key, limit)
def _get_recent_sync(self, isolation_key: str, limit: int) -> List[Dict[str, Any]]: def _get_recent_sync(self, isolation_key: str, limit: int) -> List[Dict[str, Any]]:
if self._conn is None: with self._conn_lock:
return [] if self._conn is None:
cursor = self._conn.execute( return []
"SELECT * FROM memory_history WHERE isolation_key = ? " cursor = self._conn.execute(
"ORDER BY created_at DESC LIMIT ?", "SELECT * FROM memory_history WHERE isolation_key = ? "
(isolation_key, limit), "ORDER BY created_at DESC LIMIT ?",
) (isolation_key, limit),
return [dict(row) for row in cursor.fetchall()] )
return [dict(row) for row in cursor.fetchall()]
async def count(self, isolation_key: str = "") -> int: async def count(self, isolation_key: str = "") -> int:
""" """
@@ -266,16 +274,17 @@ class HistoryStore:
return await _run_in_sqlite_pool(self._count_sync, isolation_key) return await _run_in_sqlite_pool(self._count_sync, isolation_key)
def _count_sync(self, isolation_key: str) -> int: def _count_sync(self, isolation_key: str) -> int:
if self._conn is None: with self._conn_lock:
return 0 if self._conn is None:
if isolation_key: return 0
cursor = self._conn.execute( if isolation_key:
"SELECT COUNT(*) FROM memory_history WHERE isolation_key = ?", cursor = self._conn.execute(
(isolation_key,), "SELECT COUNT(*) FROM memory_history WHERE isolation_key = ?",
) (isolation_key,),
else: )
cursor = self._conn.execute("SELECT COUNT(*) FROM memory_history") else:
return cursor.fetchone()[0] cursor = self._conn.execute("SELECT COUNT(*) FROM memory_history")
return cursor.fetchone()[0]
async def close(self) -> None: async def close(self) -> None:
"""关闭数据库连接""" """关闭数据库连接"""
@@ -284,9 +293,10 @@ class HistoryStore:
logger.info("[HistoryStore] closed") logger.info("[HistoryStore] closed")
def _close_sync(self) -> None: def _close_sync(self) -> None:
if self._conn: with self._conn_lock:
self._conn.close() if self._conn:
self._conn = None self._conn.close()
self._conn = None
# ================================================================ # ================================================================
# 内部 # 内部
@@ -339,33 +349,34 @@ class HistoryStore:
role: str = "", role: str = "",
extra: Optional[Dict[str, Any]] = None, extra: Optional[Dict[str, Any]] = None,
) -> int: ) -> int:
if self._conn is None: with self._conn_lock:
return -1 if self._conn is None:
now = _now_iso() return -1
extra_json = json.dumps(extra or {}, ensure_ascii=False) now = _now_iso()
cursor = self._conn.execute( extra_json = json.dumps(extra or {}, ensure_ascii=False)
"""INSERT INTO memory_history cursor = self._conn.execute(
(memory_id, isolation_key, event, """INSERT INTO memory_history
old_memory, new_memory, old_status, new_status, (memory_id, isolation_key, event,
change_reason, layer, actor_id, role, extra, old_memory, new_memory, old_status, new_status,
created_at, updated_at) change_reason, layer, actor_id, role, extra,
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", created_at, updated_at)
( VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
memory_id, (
isolation_key, memory_id,
event, isolation_key,
old_memory, event,
new_memory, old_memory,
old_status, new_memory,
new_status, old_status,
change_reason, new_status,
layer, change_reason,
actor_id, layer,
role, actor_id,
extra_json, role,
now, extra_json,
now, now,
), now,
) ),
self._conn.commit() )
return cursor.lastrowid self._conn.commit()
return cursor.lastrowid
+136
View File
@@ -0,0 +1,136 @@
import asyncio
import threading
import pytest
from hy_memory.config import MemoryConfig
from hy_memory.data.cache_sqlite import SqliteCache
from hy_memory.data.history_store import HistoryStore
class _BlockingCursor:
def fetchone(self):
return (1,)
def fetchall(self):
return []
class _BlockingConnection:
"""Hold execute() open so a concurrent close() can be observed."""
def __init__(self):
self.entered = threading.Event()
self.release = threading.Event()
self.closed = threading.Event()
def execute(self, *_args, **_kwargs):
self.entered.set()
if not self.release.wait(timeout=2):
raise TimeoutError("test did not release the blocking query")
if self.closed.is_set():
raise AssertionError("connection was closed while execute() was active")
return _BlockingCursor()
def close(self):
self.closed.set()
def _config(tmp_path) -> MemoryConfig:
config = MemoryConfig()
config.history.db_path = str(tmp_path / "history.db")
config.cache.db_path = str(tmp_path / "cache.db")
return config
async def _wait_for(event: threading.Event) -> None:
assert await asyncio.wait_for(asyncio.to_thread(event.wait, 1), timeout=2)
@pytest.mark.asyncio
async def test_history_close_waits_for_active_query(tmp_path):
store = HistoryStore(_config(tmp_path))
connection = _BlockingConnection()
store._conn = connection
query = asyncio.create_task(store.count())
await _wait_for(connection.entered)
closing = asyncio.create_task(store.close())
await asyncio.sleep(0.05)
assert not connection.closed.is_set()
connection.release.set()
assert await query == 1
await closing
assert connection.closed.is_set()
@pytest.mark.asyncio
async def test_cache_close_waits_for_active_query(tmp_path):
store = SqliteCache(_config(tmp_path))
connection = _BlockingConnection()
store._conn = connection
query = asyncio.create_task(store.load_metrics_range("0", "9"))
await _wait_for(connection.entered)
closing = asyncio.create_task(store.close())
await asyncio.sleep(0.05)
assert not connection.closed.is_set()
connection.release.set()
assert await query == []
await closing
assert connection.closed.is_set()
@pytest.mark.asyncio
async def test_history_concurrent_reads_and_writes_are_serialized(tmp_path):
store = HistoryStore(_config(tmp_path))
await store.initialize()
try:
inserted = await asyncio.gather(
*(
store.record_add(
memory_id=f"memory-{index}",
content="content",
layer="l2_fact",
isolation_key="user",
)
for index in range(200)
)
)
assert all(row_id > 0 for row_id in inserted)
assert len(set(inserted)) == 200
counts = await asyncio.gather(*(store.count() for _ in range(100)))
assert counts == [200] * 100
finally:
await store.close()
@pytest.mark.asyncio
async def test_cache_concurrent_pipeline_log_reads_and_writes_are_serialized(tmp_path):
store = SqliteCache(_config(tmp_path))
await store.initialize()
try:
written = await asyncio.gather(
*(
store.store_pipeline_log(
request_id="request",
user_id="user",
agent_id="agent",
step="EXTRACT",
prompt=f"prompt-{index}",
response="response",
)
for index in range(200)
)
)
assert written == [True] * 200
pages = await asyncio.gather(
*(store.get_pipeline_logs(request_id="request", limit=200) for _ in range(50))
)
assert all(len(page) == 200 for page in pages)
finally:
await store.close()