fix: harden concurrency and Redis marker safety

README.md

@@ -36,7 +36,7 @@ dp = Dispatcher()
 # Local / single process
 setup_input(dp, storage=MemoryInputStorage())
 
-# Production (multi-worker): Redis-backed wait markers
+# Production: Redis markers + TTL (futures still resolve on the owning worker)
 # redis = Redis.from_url("redis://localhost:6379/0")
 # setup_input(dp, storage=RedisInputStorage(redis, ttl=300))
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "aiogram-input"
-version = "4.2.0"
+version = "4.2.1"
 description = "Await user replies in aiogram bots with Dispatcher-scoped setup and pluggable storage."
 readme = "README.md"
 license = "MIT"

src/aiogram_input/registry.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import logging
-from asyncio import Lock
+from asyncio import InvalidStateError, Lock
 from typing import Dict, Optional
 
+from aiogram.types import Message
+
 from .types import PendingWait
 
 logger = logging.getLogger(__name__)
@@ -38,6 +40,27 @@ async def pop_if(self, chat_id: int, wait_id: str, /) -> Optional[PendingWait]:
                 return None
             return self._waits.pop(chat_id)
 
+    async def resolve_if(
+        self, chat_id: int, wait_id: str, message: Message, /
+    ) -> bool:
+        """
+        Atomically complete ``wait_id`` with ``message`` if it is still current.
+
+        Returns True only when this call consumed the wait.
+        """
+        async with self._lock:
+            wait = self._waits.get(chat_id)
+            if wait is None or wait.wait_id != wait_id:
+                return False
+            future = wait.future
+            if future.done():
+                return False
+            try:
+                future.set_result(message)
+            except InvalidStateError:
+                return False
+            return True
+
     async def contains(self, chat_id: int, /) -> bool:
         async with self._lock:
             return chat_id in self._waits
@@ -48,5 +71,8 @@ def cancel_wait(wait: PendingWait, *, chat_id: int) -> None:
         future = wait.future
         if future.done():
             return
-        future.set_result(None)
+        try:
+            future.set_result(None)
+        except InvalidStateError:
+            return
         logger.debug("[REGISTRY] Aborted leftover wait chat=%s", chat_id)

src/aiogram_input/session.py

@@ -76,16 +76,14 @@ async def feed(self, message: Message) -> bool:
             message.message_id,
         )
 
-        if not await self._storage.contains(chat_id):
-            logger.debug("[SESSION] No pending entry chat=%s", chat_id)
-            return False
-
         wait = await self._registry.get(chat_id)
         if wait is None:
-            logger.debug("[SESSION] Marker without registry wait chat=%s", chat_id)
-            await self._storage.pop(chat_id)
+            # Do NOT delete storage markers here: with Redis another worker may
+            # own the wait. Orphan Redis keys should expire via TTL.
+            logger.debug("[SESSION] No local pending wait chat=%s", chat_id)
             return False
 
+        wait_id = wait.wait_id
         if not await self._check_filter(wait.filter, message):
             filter_name = wait.filter.__class__.__name__ if wait.filter else str(None)
             logger.debug(
@@ -95,11 +93,15 @@ async def feed(self, message: Message) -> bool:
             )
             return False
 
-        if wait.future.done():
-            logger.debug("[SESSION] Future already done chat=%s", chat_id)
+        consumed = await self._registry.resolve_if(chat_id, wait_id, message)
+        if not consumed:
+            logger.debug(
+                "[SESSION] Resolve skipped stale/done wait chat=%s wait_id=%s",
+                chat_id,
+                wait_id,
+            )
             return False
 
-        wait.future.set_result(message)
         logger.debug(
             "[SESSION] Future resolved chat=%s, message_id=%s",
             chat_id,
@@ -122,7 +124,11 @@ async def _await_future(
     async def _check_filter(filter: FilterObjectType, message: Message) -> bool:
         if filter is None:
             return True
-        return await filter.call(message)
+        try:
+            return bool(await filter.call(message))
+        except Exception:
+            logger.exception("[SESSION] Filter raised; treating as reject")
+            return False
 
     async def _register_pending(
         self,

src/aiogram_input/storage/redis.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import json
+import logging
 from typing import Any, Optional
 
 from ..types import WaitRecord
 
+logger = logging.getLogger(__name__)
+
 _POP_IF_SCRIPT = """
 local raw = redis.call('GET', KEYS[1])
 if not raw then
@@ -23,8 +26,12 @@ class RedisInputStorage:
     """
     Redis-backed InputStorage for wait markers.
 
-    Futures/filters stay in-process; only ``WaitRecord`` is persisted.
-    Requires the optional dependency: ``pip install aiogram-input[redis]``.
+    Futures/filters stay in-process on the worker that called ``wait()``.
+    Redis only stores markers so TTLs and cross-process visibility of "someone
+    is waiting" work; resolving still requires the owning process (or a future
+    pub/sub bridge). Use ``ttl`` so abandoned markers cannot grow forever.
+
+    Requires: ``pip install aiogram-input[redis]``.
     """
 
     def __init__(
@@ -35,6 +42,8 @@ def __init__(
         key_prefix: str = "aiogram_input:wait:",
         ttl: Optional[int] = None,
     ) -> None:
+        if not isinstance(key_prefix, str) or not key_prefix:
+            raise ValueError("key_prefix must be a non-empty string")
         if ttl is not None and ttl <= 0:
             raise ValueError("ttl must be a positive integer or None")
         self._redis = redis
@@ -52,18 +61,29 @@ def _dumps(record: WaitRecord) -> str:
             separators=(",", ":"),
         )
 
-    @staticmethod
-    def _loads(raw: Any) -> WaitRecord:
-        if isinstance(raw, bytes):
-            raw = raw.decode()
-        data = json.loads(raw)
-        return WaitRecord(wait_id=data["wait_id"], created_at=float(data["created_at"]))
+    @classmethod
+    def _loads(cls, raw: Any) -> Optional[WaitRecord]:
+        try:
+            if isinstance(raw, bytes):
+                raw = raw.decode()
+            data = json.loads(raw)
+            wait_id = data["wait_id"]
+            created_at = float(data["created_at"])
+            if not isinstance(wait_id, str) or not wait_id:
+                raise ValueError("invalid wait_id")
+            return WaitRecord(wait_id=wait_id, created_at=created_at)
+        except (TypeError, ValueError, KeyError, json.JSONDecodeError, UnicodeError):
+            logger.warning("[REDIS] Ignoring corrupt wait marker payload")
+            return None
 
     async def get(self, chat_id: int, /) -> Optional[WaitRecord]:
         raw = await self._redis.get(self._key(chat_id))
         if raw is None:
             return None
-        return self._loads(raw)
+        record = self._loads(raw)
+        if record is None:
+            await self._redis.delete(self._key(chat_id))
+        return record
 
     async def set(self, chat_id: int, record: WaitRecord, /) -> None:
         key = self._key(chat_id)

tests/test_concurrency_hardening.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from aiogram_input.registry import WaitRegistry
+from aiogram_input.session import SessionManager
+from aiogram_input.storage import MemoryInputStorage, RedisInputStorage
+from aiogram_input.types import PendingWait, WaitRecord
+from tests.helpers import make_message
+
+
+class BoomFilter:
+    async def call(self, _message) -> bool:
+        raise RuntimeError("filter exploded")
+
+
+@pytest.mark.asyncio
+async def test_filter_exception_does_not_consume(session: SessionManager) -> None:
+    chat_id = 1
+    task = asyncio.create_task(
+        session.start_waiting(chat_id, timeout=0.2, filter=BoomFilter())  # type: ignore[arg-type]
+    )
+    await asyncio.sleep(0)
+    assert await session.feed(make_message(chat_id)) is False
+    assert await task is None
+
+
+@pytest.mark.asyncio
+async def test_resolve_if_rejects_stale_wait_id(registry: WaitRegistry) -> None:
+    loop = asyncio.get_running_loop()
+    future = loop.create_future()
+    await registry.set(
+        1, PendingWait(wait_id="current", future=future, filter=None)
+    )
+    assert await registry.resolve_if(1, "stale", make_message(1)) is False
+    assert not future.done()
+    assert await registry.resolve_if(1, "current", make_message(1)) is True
+    assert future.done()
+    assert future.result().message_id == 1
+
+
+@pytest.mark.asyncio
+async def test_concurrent_feeds_only_one_consumes(session: SessionManager) -> None:
+    chat_id = 88
+    task = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None))
+    await asyncio.sleep(0)
+    results = await asyncio.gather(
+        session.feed(make_message(chat_id, message_id=1)),
+        session.feed(make_message(chat_id, message_id=2)),
+    )
+    assert results.count(True) == 1
+    assert results.count(False) == 1
+    got = await task
+    assert got is not None
+    assert got.message_id in {1, 2}
+
+
+@pytest.mark.asyncio
+async def test_redis_corrupt_payload_is_dropped() -> None:
+    from fakeredis import FakeAsyncRedis
+
+    redis = FakeAsyncRedis(decode_responses=True)
+    storage = RedisInputStorage(redis, key_prefix="bad:")
+    await redis.set("bad:1", "{not-json")
+    assert await storage.get(1) is None
+    assert await redis.exists("bad:1") == 0
+    await redis.aclose()
+
+
+@pytest.mark.asyncio
+async def test_redis_rejects_empty_key_prefix() -> None:
+    from fakeredis import FakeAsyncRedis
+
+    with pytest.raises(ValueError, match="key_prefix"):
+        RedisInputStorage(FakeAsyncRedis(), key_prefix="")
+
+
+@pytest.mark.asyncio
+async def test_feed_after_overwrite_goes_to_latest_wait(
+    session: SessionManager,
+) -> None:
+    chat_id = 77
+    first = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None))
+    await asyncio.sleep(0)
+    second = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None))
+    assert await first is None
+    msg = make_message(chat_id, message_id=9)
+    assert await session.feed(msg) is True
+    assert await second is msg

tests/test_edge_cases.py

@@ -36,12 +36,13 @@ async def handler(event, data):
 
 
 @pytest.mark.asyncio
-async def test_orphan_storage_marker_is_cleaned_on_feed(
+async def test_storage_marker_without_local_wait_is_not_deleted(
     session: SessionManager, storage: MemoryInputStorage
 ) -> None:
+    """Redis multi-worker safety: never delete another process's marker."""
     await storage.set(55, WaitRecord(wait_id="orphan", created_at=1.0))
     assert await session.feed(make_message(55)) is False
-    assert await storage.contains(55) is False
+    assert await storage.contains(55) is True
 
 
 @pytest.mark.asyncio
@@ -55,6 +56,8 @@ async def test_feed_when_future_already_done(session: SessionManager) -> None:
     )
     await session._storage.set(chat_id, WaitRecord(wait_id="x", created_at=1.0))
     assert await session.feed(make_message(chat_id)) is False
+    # Marker may remain until owning cleanup/TTL; feed must not crash.
+    assert await session._storage.contains(chat_id) is True
 
 
 @pytest.mark.asyncio

tests/test_redis_storage.py

@@ -78,3 +78,9 @@ async def test_redis_loads_bytes_payload() -> None:
     record = WaitRecord(wait_id="b", created_at=3.0)
     raw = RedisInputStorage._dumps(record).encode()
     assert RedisInputStorage._loads(raw) == record
+
+
+@pytest.mark.asyncio
+async def test_redis_loads_invalid_payload_returns_none() -> None:
+    assert RedisInputStorage._loads(b"nope") is None
+    assert RedisInputStorage._loads('{"wait_id":""}') is None