From Manual SQL to pytest+TiDB: Cutting AI Agent Memory Regression from 2 Hours to 5 Minutes

At 3 AM, our ops colleague frantically @-ed me in the group chat: the online AI agent had suddenly started "forgetting" — a user preference it just learned a moment ago would be gone the next. Digging in for half an hour, I found the culprit: concurrent memory writes weren't locking the version number properly, and the version field was being silently overwritten. What drove me crazy was that we had verified this exact scenario before the release by "manually running SQL", and everything looked fine — because human eyes simply can't spot that one missing UPDATE row.

This is the brutal truth about memory storage consistency: it always "looks fine" during manual checks, and then concurrency punches you in the face the moment you hit production.

Breaking Down the Problem

An AI agent's memory is not a simple key-value store. It's versioned, incremental storage: each (session_id, key) pair has a linear history. Every write must read the latest version, increment it by 1, and then write — otherwise you get overwrites. TiDB is MySQL-compatible, but under the default optimistic transaction model, two concurrent requests may read the same version, both attempt an UPDATE, and only one succeeds while the other gets silently discarded. That was exactly the root cause on that night.

What's wrong with the usual approaches?

• Manual SQL verification: Open two terminal windows in the test environment, manually fire UPDATEs, then SELECT and compare. It's slow, irreproducible, and you can only realistically watch two or three concurrent operations at once. It can never simulate the thousands of QPS you see in production.
• Unit tests with a mocked database: Replace TiDB with an in-memory SQLite store. The transaction isolation is completely different — you're testing nothing.
• Data comparison scripts: Write a Python script that runs SELECT twice and compares results. It can spot inconsistencies, but running it manually for every release takes at least two hours. Who can live with that?

We needed an automated solution that could simulate real concurrency, assert consistency, and run blazingly fast. That's when I turned my attention to the combination of pytest and TiDB.

Design Decisions

Why these choices?

• Why pytest? Its fixture mechanism is a natural fit for managing database connections and cleanup. pytest-xdist can distribute concurrency tests across multiple processes to simulate high concurrency. pytest-ordering lets you control execution order. This is flexibility that unittest simply can't offer.
• Why TiDB? The agent's memory store needs horizontal scalability. TiDB is a distributed database that provides the same strong consistency guarantees as MySQL (in pessimistic transaction mode), and you can spin up a single-node cluster for CI using tiup — no difference from local development.
• Why not other databases? CockroachDB's syntax differences add migration cost. MySQL Group Replication is heavy and not truly distributed. PostgreSQL is powerful, but the ecosystem tooling around it lacks the MySQL-compatibility convenience that TiDB offers. The most important point: production runs on TiDB; the test environment must match production. That's an iron rule.

The architecture is straightforward: pytest orchestrates the test logic and concurrency, TiDB provides a real distributed transaction environment, and pymysql is the driver. Each test case is an independent "memory operation scenario" that verifies eventual consistency through concurrent execution.

Core Implementation

1. Setting Up the Test Harness: Connection Management, Table Creation, and Cleanup

This code ensures every test starts from a clean state and leaves no garbage behind. I use a pytest.fixture to automatically create the table before a test and drop it afterward.

# conftest.py
import pytest
import pymysql

@pytest.fixture(scope="function")
def db_conn():
    # 连接TiDB——这里可以是本地tiup集群，也可以是TiDB Cloud Serverless
    conn = pymysql.connect(
        host="127.0.0.1",
        port=4000,
        user="root",
        password="",
        database="test",
        charset="utf8mb4",
        autocommit=False          # 显式控制事务，方便模拟并发
    )
    # 建表：核心是version字段，每次更新必须原子递增
    with conn.cursor() as cur:
        cur.execute("""
            CREATE TABLE IF NOT EXISTS agent_memory (
                id BIGINT AUTO_RANDOM PRIMARY KEY,
                session_id VARCHAR(64) NOT NULL,
                mem_key VARCHAR(128) NOT NULL,
                mem_value TEXT,
                version BIGINT NOT NULL DEFAULT 0,
                UNIQUE KEY uk_session_key (session_id, mem_key)
            )
        """)
        conn.commit()
    yield conn
    # 清理：直接删表，干净利落
    with conn.cursor() as cur:
        cur.execute("DROP TABLE IF EXISTS agent_memory")
        conn.commit()
    conn.close()

2. The First Hardcore Test Case: Concurrent Updates Must Be Strictly Sequential

This test verifies that when two requests update the same memory key almost simultaneously, the version must increase by exactly 2 and no write must be lost. I simulate concurrency using threading because TiDB's pessimistic transactions will queue on row locks — perfect for validating the locking mechanism.

# test_memory_consistency.py
import threading
import pytest
import pymysql

@pytest.fixture
def init_row(db_conn):
    """预先插入一条记忆记录"""
    with db_conn.cursor() as cur:
        cur.execute(
            "INSERT INTO agent_memory (session_id, mem_key, mem_value, version) "
            "VALUES ('sess1', 'prefer', 'old', 1)"
        )
        db_conn.commit()
    yield
    # 不用手动清理，db_conn fixture会删表

def update_memory(conn, new_value, result_holder):
    """并发线程执行的更新逻辑——标准乐观锁写法，但依赖于TiDB的悲观锁保证不丢失"""
    try:

Comments

[Alex Shev]

Turning memory regression into pytest is the right move. Agent memory should be treated like product behavior: fixtures, expected recall, stale-data cases, and a repeatable failure when the memory layer drifts.