We’ve Been Testing LangChain Memory Rollbacks Wrong for 6 Months

python dev.to

It was 3 a.m. when a message from a customer jolted me awake: “Your support bot is amnesic again — after rolling back a conversation, it started mixing up things from last month.” I pulled up the monitoring dashboards, and the logs confirmed it: during a session rollback, ConversationBufferMemory was leaking irrelevant context into a freshly branched conversation. Classic memory contamination.

What broke me? The “rollback test suite” we’d been running weekly for half a year had passed every single time. We had been testing memory persistence with a completely wrong approach, and the real bugs — the forgetting, the cross‑branch leakage — only surfaced once we put the whole chat pipeline under a true end‑to‑end rollback with Playwright + pytest, capturing everything that hides inside UI interactions and async storage.

Why memory rollbacks are a deep‑water zone

LLM memory stores like LangChain’s ConversationBufferMemory and ConversationSummaryMemory almost universally support rollback — the ability for a user to jump back to an earlier message and restart the conversation from that point. This is a hard requirement for chatbots in customer support, education, and game NPCs. Implementing it, however, means solving two core problems:

  • Forgetting: after rolling back, earlier context gets discarded incorrectly, leaving the model without necessary history.
  • Contamination: after rolling back, traces from the branch that should have been pruned remain, so the model starts mixing up content from different conversation paths.

The typical testing approach is to call memory.load_memory_variables() in a unit test or fire a few messages with Postman and compare the responses by eye. But in production, memory persistence almost always depends on:

  1. session state that spans continuous frontend interactions,
  2. asynchronous message writes and caching layers,
  3. concurrent reads of the memory store by multiple microservices.

Manual testing simply cannot cover real‑world timing, while unit tests are too idealised — they treat memory as a pure function and ignore the “dirty” parts: frontend event loops, network latency, API idempotency. When rollback behaviour involves a multi‑step dance between front and backend, a single manual regression run can easily take over ten minutes and still might not reproduce the bug.

So the root cause wasn’t a faulty memory implementation. It was a disconnect between the test method and the real execution path. What we lacked was an end‑to‑end rollback test that could faithfully simulate real user actions and automatically assert on the memory state.

Solution design: Playwright + pytest as a “time traveller” for memory

For end‑to‑end rollback testing we looked at a few options:

  • Pure API testing (requests + pytest): quick to call endpoints, but it cannot simulate page navigations, WebSocket reconnections, or other frontend behaviour, making it easy to miss issues with frontend state synchronisation.
  • Selenium: the veteran tool, but its async‑waiting mechanisms are heavy, and its support for modern SPAs and WebSockets is less natural than Playwright’s built‑in capabilities.
  • Playwright: supports multiple browsers, native auto‑waiting, network interception, and a trace viewer. It can perfectly mimic a real user clicking, typing, and rolling back in the chat UI. Combined with pytest, we also get fixture reuse, parametrisation, and parallel execution for free.

The final architecture is lightweight: pytest manages the test lifecycle, Playwright drives the browser through a complete conversation flow (send messages, click the rollback button, inspect the history list), and the existing FastAPI chat service on the backend — which already integrates LangChain’s ConversationBufferMemory — handles the API layer. After each test case we reset the memory store via an API call to ensure isolation.

Why didn’t we use LangChain’s built‑in testing tools? Because the official docs only provide examples for unit‑testing memory — there is no end‑to‑end solution, and they can’t cover the UI layer at all. What we wanted to test was the memory rollback “through the eyes of the user”, not just the code logic.

Core implementation: making the rollback test a reusable playbook

Below is the full test pipeline. Let’s start with a minimal backend so you can reproduce the setup. This code provides a chat API that supports rollback.

# app.py
import uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain.memory import ConversationBufferMemory
from langchain.schema import HumanMessage, AIMessage

app = FastAPI()
sessions: dict[str, ConversationBufferMemory] = {}

class ChatRequest(BaseModel):
    session_id: str
    message: str

class RollbackRequest(BaseModel):
    session_id: str
    target_index: int  # 回滚到第几条消息之后(从0开始)

@app.post("/chat")
async def chat(req: ChatRequest):
    mem = sessions.get(req.session_id)
    if not mem:
        mem = ConversationBufferMemory(return_messages=True)
        sessions[req.session_id] = mem

    # 模拟大模型回复:简单回显用户输入并记录
    mem.chat_memory.add_user_message(req.message)
    response = f"Echo: {req.message}"
    mem.chat_memory.add_ai_message(response)
    return {"reply": response}

@app.post("/rollback")
async def rollback(req: RollbackRequest):
    mem = sessions.get(req.session_id)
    if not mem:
        raise HTTPException(status_code=404, detail="Session not found")
    messages = mem.chat_memory.messages
    if req.target_index < 0 or req.target_index * 2 > len(messages):
        raise HTTPException(status_code=400, detail="Invalid index")
    # 截断消息列表到指定轮次:保留前 target_index*2 条消息(一轮 = user + ai)
    mem.chat_memory.messages = messages[: req.target_index * 2]
    return {"status": "rolled back", "remaining_turns": len(messages) // 2}

@app.post("/reset/{session_id}")
async def reset(session_id: str):
    if session_id in sessions:
        del sessions[session_id]
    return {"status": "reset"}
Enter fullscreen mode Exit fullscreen mode

The corresponding Playwright test uses pytest fixtures to control a browser, walk through a multi‑turn conversation, trigger the rollback UI action, and then verify that only the correct messages remain visible.

# test_chat_rollback.py
import pytest
from playwright.sync_api import Page, expect

@pytest.fixture(scope="module")
def browser_context(browser):
    context = browser.new_context()
    yield context
    context.close()

@pytest.fixture
def page(browser_context):
    p = browser_context.new_page()
    yield p
    p.close()

def test_rollback_does_not_contaminate(page: Page):
    # 打开聊天页面
    page.goto("http://localhost:3000/chat")
    session_id = page.evaluate("window.__SESSION_ID__")

    # 第一轮对话
    page.fill('[data-testid="chat-input"]', "Hello, my name is Alex")
    page.click('[data-testid="send-button"]')
    expect(page.locator('[data-testid="message"]').last).to_contain_text("Echo: Hello, my name is Alex")

    # 第二轮对话
    page.fill('[data-testid="chat-input"]', "I live in Paris")
    page.click('[data-testid="send-button"]')
    expect(page.locator('[data-testid="message"]').last).to_contain_text("Echo: I live in Paris")

    # 回滚到第0轮之后(即保留第一轮,删除第二轮)
    page.click(f'[data-testid="rollback-0"]')
    page.wait_for_timeout(500)  # 等待UI更新完成

    # 断言:第一轮消息仍然存在
    messages = page.locator('[data-testid="message"]').all_text_contents()
    assert any("Hello, my name is Alex" in msg for msg in messages)
    # 断言:第二轮消息已消失
    assert not any("I live in Paris" in msg for msg in messages)

    # 额外断言:后端记忆数量应为2条(1 user + 1 ai)
    import requests
    resp = requests.get(f"http://localhost:8000/memory/{session_id}")
    assert len(resp.json()["messages"]) == 2

    # 清理:重置session
    requests.post(f"http://localhost:8000/reset/{session_id}")
Enter fullscreen mode Exit fullscreen mode

The test above treats the system as a black box from the user’s perspective. It does not touch ConversationBufferMemory directly — instead it drives the UI, sends HTTP requests where necessary for finer‑grained assertions, and ensures that what the user sees matches what’s stored on the backend. This is the exact opposite of a pure unit test; it catches the kind of contamination that only appears when the frontend’s session management, the backend’s in‑memory store, and the asynchronous update cycle interact in unintended ways.

What we found after running the real tests

The first run immediately exposed a contamination bug: when the frontend optimistically updates the UI after a rollback, it sometimes re‑renders stale cached messages from a previous branch, while the backend has already applied the truncation correctly. A unit test of the memory store alone would never see this, because the divergence only exists in the gap between the frontend cache and the backend state.

Another discovery was that our message indexing depended on the client’s local timestamp ordering, which drifted when messages arrived out of order during WebSocket reconnection — the exact scenario that a rollback often triggers. The Playwright test reproduced this reliably because it could simulate rapid back‑and‑forth navigation and inspect the DOM after every step.

The most satisfying moment? After we integrated these tests into our CI pipeline, the “unexplainable amnesia” tickets from customers dropped to zero.

Lessons for anyone building with LLM memory

  1. Don’t trust unit tests for stateful, interactive memory. If your product involves a chat UI and rollback, you must test the full chain: browser → frontend state → WebSocket / REST → backend memory store.
  2. Treat rollback contamination as a UI‑visible problem first. The user sees it as inconsistent history, not as a corrupted messages list.
  3. Playwright + pytest is the secret sauce. It gives you the observability to trace exactly what the user saw, and the parametrisation to turn long, flaky manual regression runs into deterministic, fast‑feedback CI checks.

We spent six months running the wrong tests and sleeping peacefully. The moment we ran the right ones, the bugs jumped out of the screen. If you’re building chatbot memory today, don’t wait for a 3 a.m. wake‑up call — write a real end‑to‑end rollback test first.

Source: dev.to

arrow_back Back to Tutorials