LLM记忆系统踩坑实录:一个缓存丢失的bug让我排查了3小时,最后用pytest快照测试彻底解决

4 阅读1分钟

凌晨两点,我被报警电话震醒——线上Agent连续三次对话都“失忆”了,之前说好的上下文全丢,用户疯狂投诉。我眯着眼睛打开日志,发现记忆管理模块的rollback方法被一个看似无害的代码重构给搞坏了:回滚不仅撤销了错误操作,还把整个会话历史一并干掉了。最要命的是,这个问题在我们现有的单元测试里完全没暴露——因为测试每次都从空数据库开始,根本覆盖不了“脏数据回滚到上一个快照”这种跨会话场景。我花了3小时排查、手动模拟状态,才定位到根因。那一刻我才意识到:我们缺的不是测试,是能把“记忆状态”整个拍下来的快照测试。


问题拆解

我们的LLM记忆系统基于SQLite做本地持久化,每个会话维护一张表,保存对话轮次、向量摘要和工具调用记录。关键操作有两个:

  • save_snapshot(session_id):把当前会话的完整状态序列化到snapshots表,生成一个可回滚的检查点。
  • rollback_to_snapshot(session_id, snapshot_id):发生异常时,根据快照重建会话表,丢弃之后的所有变更。

这套机制平时跑得好好的,但那次重构我改动了回滚逻辑中的事务边界,导致rollback执行后,conversation表被重建,但snapshots表自身也被意外清空——下一次回滚就再也找不到更早的检查点了。

常规单元测试为什么抓不住这个bug?因为测试套路是:

def test_rollback():
    db = create_in_memory_db()
    db.save_snapshot("s1")
    db.rollback_to_snapshot("s1", ...)
    assert db.get_conversation("s1") == expected

这一切都发生在一个进程、一个临时数据库里。但线上问题是:进程A保存快照后退出,进程B重新打开同一数据库文件并执行回滚。文件级的持久化状态、WAL日志的合并、甚至不同连接对snapshots表的可见性,这些统统没测到。说白了,我们测了“逻辑”,却没测“存储”。


方案设计

我决定引入**快照测试(snapshot testing)**的思想,但不用常见的文本快照,而是把SQLite数据库文件本身当作不可变快照。

选型对比:

  • pytest-snapshot:只做文本/JSON快照,不适合二进制或复杂状态比对。
  • 直接用pytest的tmp_path+手动比对:灵活,但每次自己写比对逻辑容易遗漏字段。
  • 基于文件hash + SQLite内对比:先对数据库文件计算sha256作为整体快照,同时可选导出关键表的内容做可读差异对比。这样做既能快速发现“哪里变了”,又能保留细粒度调试能力。

架构思路:pytest的conftest.py提供一个snapshot_db fixture,它负责:

  1. 在测试开始前,检查是否存在基线快照文件(比如tests/snapshots/memory_test.sqlite)。
  2. 如果没有,则自动生成(--snapshot-update模式),测试直接pass。
  3. 如果有,则执行测试操作后,对被测数据库文件做sha256,并与基线文件的sha256对比。如果不一致,测试失败并输出差异提示。

这样,我们的测试就真正模拟了“跨进程、跨连接”的持久化效果——每个测试用例拿到的都是一个独立拷贝的数据库文件,执行完操作后比较整个文件状态是否与预期完全一致。


核心实现

1. 先搭一个可持久化的记忆管理器

这段代码解决“我们到底要测什么”的问题。MemoryManager封装了SQLite的连接、快照保存和回滚,是线上真实使用的简化版。

# memory_manager.py
import sqlite3
import uuid
from datetime import datetime, timezone

class MemoryManager:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._init_tables()

    def _get_conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path)
        conn.execute("PRAGMA journal_mode=WAL")
        conn.row_factory = sqlite3.Row
        return conn

    def _init_tables(self):
        with self._get_conn() as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS conversations (
                    session_id TEXT NOT NULL,
                    turn INTEGER NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    PRIMARY KEY (session_id, turn)
                );
                CREATE TABLE IF NOT EXISTS snapshots (
                    snapshot_id TEXT PRIMARY KEY,
                    session_id TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    state_json TEXT NOT NULL
                );
            """)

    def add_message(self, session_id: str, role: str, content: str):
        with self._get_conn() as conn:
            turn = conn.execute(
                "SELECT COALESCE(MAX(turn), 0) + 1 FROM conversations WHERE session_id = ?",
                (session_id,)
            ).fetchone()[0]
            conn.execute(
                "INSERT INTO conversations VALUES (?, ?, ?, ?)",
                (session_id, turn, role, content)
            )

    def save_snapshot(self, session_id: str) -> str:
        snapshot_id = str(uuid.uuid4())
        with self._get_conn() as conn:
            rows = conn.execute(
                "SELECT turn, role, content FROM conversations WHERE session_id = ? ORDER BY turn",
                (session_id,)
            ).fetchall()
            import json
            state = [dict(r) for r in rows]
            conn.execute(
                "INSERT INTO snapshots VALUES (?, ?, ?, ?)",
                (snapshot_id, session_id, datetime.now(timezone.utc).isoformat(), json.dumps(state))
            )
        return snapshot_id

    def rollback_to_snapshot(self, session_id: str, snapshot_id: str):
        with self._get_conn() as conn:
            row = conn.execute(
                "SELECT state_json FROM snapshots WHERE snapshot_id = ?",
                (snapshot_id,)
            ).fetchone()
            if not row:
                raise ValueError("Snapshot not found")
            import json
            state = json.loads(row["state_json"])
            conn.execute("DELETE FROM conversations WHERE session_id = ?", (session_id,))
            for msg in state:
                conn.execute(
                    "INSERT INTO conversations VALUES (?, ?, ?, ?)",
                    (session_id, msg["turn"], msg["role"], msg["content"])
                )

2. 快照fixture:将数据库变成可验证的“照片”

这段代码解决“怎么自动比对整个数据库文件”的问题。fixture会复制一个干净的模板数据库,测试操作完成后,对最终数据库文件取sha256,与基线快照文件的sha256比较。

# conftest.py
import pytest
import shutil
import hashlib
import os
from pathlib import Path

SNAPSHOT_DIR = Path(__file__).parent / "snapshots"

def _db_checksum(path: str) -> str:
    with open(path, "rb") as f:
        return hashlib.sha256(f.read()).hexdigest()

@pytest.fixture
def snapshot_db(request, tmp_path):
    """
    返回一个 MemoryManager 实例,其背后是一个可落盘的数据库文件。
    测试结束后,自动比对数据库文件与基线快照。
    """
    # 基线快照路径:测试函数名+.sqlite
    snapshot_name = request.node.name + ".sqlite"
    snapshot_path = SNAPSHOT_DIR / snapshot_name

    # 在临时目录创建被测数据库(先复制基线,如果存在)
    db_tmp = tmp_path / "test.db"
    if snapshot_path.exists():
        shutil.copy(snapshot_path, db_tmp)
    else:
        # 首次运行:创建空数据库,后续由 --snapshot-update 模式生成基线
        pass

    manager = MemoryManager(str(db_tmp))
    yield manager

    # 检查是否需要更新快照
    update_snapshots = request.config.getoption("--snapshot-update", default=False)
    manager._get_conn().close()  # 确保 WAL 刷盘

    if update_snapshots:
        SNAPSHOT_DIR.mkdir(parents=True, exist_ok=True)
        shutil.copy(db_tmp, snapshot_path)
    else:
        if not snapshot_path.exists():
            pytest.fail(f"Snapshot file {snapshot_path} missing. Run with --snapshot-update to create.")
        expected_checksum = _db_checksum(str(snapshot_path))
        actual_checksum = _db_checksum(str(db_tmp))
        if actual_checksum != expected_checksum:
            pytest.fail(
                f"Database snapshot mismatch.\n"
                f"Expected: {expected_checksum}\n"
                f"Actual:   {actual_checksum}\n"
                f"Tip: run with --snapshot-update to regenerate."
            )

3. 写测试用例:验证保存后回滚不会丢数据

这段代码模拟了“存快照→加新消息→回滚”的关键路径,并依靠fixture自动核对数据库完整性。

# test_memory.py
import pytest

def test_rollback_preserves_earlier_snapshots(snapshot_db):
    session = "sess_01"
    # 初始对话
    snapshot_db.add_message(session, "user", "hello")
    snapshot_db.add_message(session, "assistant", "hi there")
    snap_id = snapshot_db.save_snapshot(session)

    # 再次对话并产生新快照
    snapshot_db.add_message(session, "user", "how are you?")
    snapshot_db.add_message(session, "assistant", "I'm fine")
    snap_id2 = snapshot_db.save_snapshot(session)

    # 回滚到第一个快照
    snapshot_db.rollback_to_snapshot(session, snap_id)

    # 此时数据库状态应等于第一次快照后的状态
    # 快照fixture会在测试结束后做全库比对,如果意外删除了 snapshots 表内容,文件checksum必变
    # 我们还可以做轻量断言辅助定位
    rows = snapshot_db._get_conn().execute(
        "SELECT turn, content FROM conversations WHERE session_id=? ORDER BY turn", (session,)
    ).fetchall()
    assert len(rows) == 2
    assert rows[1]["content"] == "hi there"

此刻,任何破坏回滚逻辑的修改——比如错误地清空了snapshots表——都会被文件级的checksum比对抓出来,因为清空操作改变了数据库页,即使最终会话数据看起来“对”了,整个文件也不再匹配基线。


踩坑记录

坑1:WAL模式下的“影子数据库”

现象:第一次跑快照测试,checksum总是不一样,即使逻辑完全正确。打开基线文件对比,大小有时差几十KB。原因:SQLite在WAL模式下,主数据库文件不会立即写入最新数据,而是先写-wal文件,只有checkpoint之后才会合并。我们的fixture在测试结束时直接对test.db读hash,可能读到一个尚未合并WAL的状态。

解决:在finish阶段,显式关闭连接(manager._get_conn().close())并在关闭前执行PRAGMA wal_checkpoint(TRUNCATE)。但最简单的方式是关闭所有连接后,SQLite会自动做checkpoint。我们的fixture在关闭连接后才计算hash,保证了文件完整性。这是官方文档没明说、实际调试时才会发现的细节。

坑2:快照文件没进版本控制,CI全红

现象:本地跑全绿,推到CI上直接挂,报Snapshot file missing。因为我们团队有人忘记git add新生成的快照文件,而CI环境从不运行--snapshot-update解决:在CI脚本里加了一个防御性检查——如果tests/snapshots/下有未跟踪文件,直接失败并提示开发者提交。同时在README里写清楚:任何修改记忆管理模块的PR,必须附带更新后的快照文件。这也倒逼团队养成了“改持久化先看快照”的习惯。


效果验证

以前手动回归整个记忆回滚流程,需要分3个场景、来回构造SQLite文件,最快也要30分钟,还容易漏测。现在:

  • 测试覆盖场景数:1个 → 7个(正常回滚、回滚到空快照、回滚不存在的id、并发保存、WAL积压后回滚等)
  • 单次回归耗时:30分钟 → 12秒
  • Bug发现率:线上记忆丢失事件后再无同类故障,两次重构均被快照测试挡在PR阶段

快照比对让“持久化是否变脏”变成二进制判断,再也不用靠人肉看日志猜状态了。


直接可用的代码

# 一行命令更新全部快照基线,开发者本地执行
pytest --snapshot-update

把这套conftest.pySNAPSHOT_DIR放入你的项目,接入任何依赖文件状态的模块,立刻获得跨进程的持久化验证能力。


#Python #LLM #测试 #后端 #SQLite

关于作者
一个专啃后端硬骨头的实战派开发者,擅长把内存、存储、并发这些“玄学”问题变成代码就能验证的确定性方案。
GitHub: github.com/baofugege — 本项目的完整快照测试框架模板即将开源。
Sponsor: github.com/sponsors/ba… — 如果这篇文章帮你省下3小时排查时间,请我喝杯咖啡。
提供服务:Python后端性能优化 / 工具定制 / 技术咨询,联系 Telegram @baofugege