From 84c2b21b96f07070c65d36386e6758e0e8a88d22 Mon Sep 17 00:00:00 2001 From: "zoz.eth" Date: Sun, 22 Mar 2026 01:57:04 +0000 Subject: [PATCH] feat: add RedisStore for production multi-instance deployments --- pyproject.toml | 1 + src/mpp/stores/__init__.py | 5 +++ src/mpp/stores/redis.py | 60 +++++++++++++++++++++++++++++++++ tests/test_stores_redis.py | 68 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 134 insertions(+) create mode 100644 src/mpp/stores/__init__.py create mode 100644 src/mpp/stores/redis.py create mode 100644 tests/test_stores_redis.py diff --git a/pyproject.toml b/pyproject.toml index 275151f..c2b59e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ tempo = [ ] server = ["pydantic>=2.0", "python-dotenv>=1.0"] mcp = ["mcp>=1.1.0"] +redis = ["redis>=5.0"] dev = [ "pytest>=8.0", "pytest-asyncio>=0.24", diff --git a/src/mpp/stores/__init__.py b/src/mpp/stores/__init__.py new file mode 100644 index 0000000..f4b1085 --- /dev/null +++ b/src/mpp/stores/__init__.py @@ -0,0 +1,5 @@ +"""Store implementations for MPP servers.""" + +from mpp.stores.redis import RedisStore + +__all__ = ["RedisStore"] diff --git a/src/mpp/stores/redis.py b/src/mpp/stores/redis.py new file mode 100644 index 0000000..db75465 --- /dev/null +++ b/src/mpp/stores/redis.py @@ -0,0 +1,60 @@ +"""Redis-backed store for MPP challenge replay protection.""" +from __future__ import annotations + +from typing import Any + +_SETNX_SCRIPT = """ +local set = redis.call('SET', KEYS[1], ARGV[1], 'NX', 'EX', ARGV[2]) +if set then return 1 else return 0 end +""" + + +class RedisStore: + """Production-ready Redis store with atomic SETNX and automatic TTL. + + Suitable for multi-instance deployments where MemoryStore would allow + replay attacks across processes. + + Example:: + + from redis.asyncio import from_url + from mpp.stores.redis import RedisStore + + store = RedisStore(await from_url("redis://localhost:6379")) + + Args: + client: An async Redis client (e.g. ``redis.asyncio.Redis``). + ttl_seconds: How long challenges remain valid. Defaults to 300 (5 min). + key_prefix: Namespace prefix for all keys. Defaults to ``"mpp:"``. + """ + + def __init__( + self, + client: Any, + *, + ttl_seconds: int = 300, + key_prefix: str = "mpp:", + ) -> None: + self._redis = client + self._ttl = ttl_seconds + self._prefix = key_prefix + + def _key(self, key: str) -> str: + return f"{self._prefix}{key}" + + async def get(self, key: str) -> Any | None: + value = await self._redis.get(self._key(key)) + return value # bytes or str depending on decode_responses + + async def put(self, key: str, value: Any) -> None: + await self._redis.set(self._key(key), value, ex=self._ttl) + + async def delete(self, key: str) -> None: + await self._redis.delete(self._key(key)) + + async def put_if_absent(self, key: str, value: Any) -> bool: + """Atomic SETNX with TTL — maps to Redis SET NX EX.""" + result = await self._redis.set( + self._key(key), value, nx=True, ex=self._ttl + ) + return result is not None diff --git a/tests/test_stores_redis.py b/tests/test_stores_redis.py new file mode 100644 index 0000000..4bf3c08 --- /dev/null +++ b/tests/test_stores_redis.py @@ -0,0 +1,68 @@ +"""Tests for RedisStore.""" +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from mpp.stores.redis import RedisStore + + +@pytest.fixture +def redis_client() -> MagicMock: + client = MagicMock() + client.get = AsyncMock(return_value=None) + client.set = AsyncMock(return_value=True) + client.delete = AsyncMock(return_value=1) + return client + + +@pytest.fixture +def store(redis_client: MagicMock) -> RedisStore: + return RedisStore(redis_client, ttl_seconds=60, key_prefix="test:") + + +async def test_put(store: RedisStore, redis_client: MagicMock) -> None: + await store.put("challenge-1", "value-1") + redis_client.set.assert_awaited_once_with("test:challenge-1", "value-1", ex=60) + + +async def test_get_returns_value(store: RedisStore, redis_client: MagicMock) -> None: + redis_client.get.return_value = b"stored-value" + result = await store.get("challenge-1") + redis_client.get.assert_awaited_once_with("test:challenge-1") + assert result == b"stored-value" + + +async def test_get_returns_none_when_missing(store: RedisStore, redis_client: MagicMock) -> None: + redis_client.get.return_value = None + result = await store.get("missing") + assert result is None + + +async def test_delete(store: RedisStore, redis_client: MagicMock) -> None: + await store.delete("challenge-1") + redis_client.delete.assert_awaited_once_with("test:challenge-1") + + +async def test_put_if_absent_returns_true_when_key_absent( + store: RedisStore, redis_client: MagicMock +) -> None: + redis_client.set.return_value = True # Redis returns OK when SET NX succeeds + result = await store.put_if_absent("challenge-1", "value-1") + redis_client.set.assert_awaited_once_with("test:challenge-1", "value-1", nx=True, ex=60) + assert result is True + + +async def test_put_if_absent_returns_false_when_key_exists( + store: RedisStore, redis_client: MagicMock +) -> None: + redis_client.set.return_value = None # Redis returns None when SET NX fails + result = await store.put_if_absent("challenge-1", "value-1") + assert result is False + + +async def test_key_prefix(redis_client: MagicMock) -> None: + store = RedisStore(redis_client, key_prefix="mpp:") + await store.get("foo") + redis_client.get.assert_awaited_once_with("mpp:foo")