DEV Community

BAOFUFAN
BAOFUFAN

Posted on

Debugging AI Agent Context Pollution with pytest + Redis: 3 Cross-Talk Bugs in 6 Hours

At 2 a.m., a screenshot landed in the user group: the AI support agent replied to a ticket and leaked another user's phone number and shipping address. My first instinct was model hallucination, but digging into Redis revealed that two sessions with different user_id values shared the same session_id as the key — the context was cross-wired. This wasn't a model problem; the memory layer was running loose. I later built a consistency test suite with pytest + Redis to finally kill this class of bugs.

Breaking Down the Problem

Memory storage for AI agents is different from a normal cache: losing a cache only makes things slower, but mixing up memories directly leaks privacy and derails conversations. Our team stores the agent's short-term memory in Redis, relying on it to keep multi-turn conversations coherent. After launch, we occasionally saw user A seeing user B's history. The root cause wasn't the model — it was three issues in the memory layer: key design lacking a user dimension, messy TTL refresh logic, and inconsistent serialization protocols. Why weren't regular approaches enough? Code review and manual testing can't catch this because cross-talk only appears under concurrency plus edge-case combinations; mocking Redis can't verify real TTL behavior or atomicity. In the end, I decided to treat the memory layer as infrastructure and write consistency tests with pytest + Redis.

Solution Design

The reason to choose pytest + Redis instead of unit-test mocks is simple: Redis's real behavior — TTL calculation, WATCH optimistic locking, and binary differences after serialization — cannot be simulated by mocks. Pytest fixtures can give each test a clean database, and a real Redis can be spun up with Docker. In the design, we wrap memory reads and writes into MemoryStore, then write tests against three invariants: key isolation, TTL/version compatibility, and concurrent append consistency. Why not test directly against production? Too dangerous. Why not only write unit tests? Because they treat Redis like a plain dict and hide TTL and concurrency problems.

Core Implementation

This code solves the test isolation problem: each test case gets an independent Redis db so leftover data doesn't pollute assertions.

# conftest.py
import pytest
import redis

@pytest.fixture
def redis_client():
    """连接本地 Redis,测试前清空当前 db,保证隔离"""
    client = redis.Redis(host="localhost", port=6379, db=15, decode_responses=True)
    client.flushdb()
    yield client
    client.flushdb()
    client.close()
Enter fullscreen mode Exit fullscreen mode

This class encapsulates production-grade Redis memory reads and writes, including key namespacing, JSON versioning, SETEX TTL, and concurrency-safe WATCH appends. Key lines are commented.


python
# memory_store.py
import json
from typing import Any, Dict, List, Optional

import redis
from redis.exceptions import WatchError

class MemoryStore:
    """AI Agent 短期记忆存储,负责 key 隔离、TTL 和版本化序列化"""

    def __init__(self, client: redis.Redis, ttl: int = 3600):
        self.client = client
        self.ttl = ttl
        self.version = 1  # 序列化
Enter fullscreen mode Exit fullscreen mode

Top comments (0)