webagent-protocol 0.2.0__py3-none-any.whl

wap/storage/__init__.py

@@ -0,0 +1,30 @@
+"""Pluggable state storage for WAP servers (in-memory by default, Redis for multi-worker deployments)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .base import LockTimeout, StateStore
+from .memory import MemoryStore
+from .meters import MeterTable, RateLimitDecision, SlidingWindowCounter, TokenBucket
+
+
+def __getattr__(name: str) -> Any:
+    # RedisStore needs the optional ``redis`` package; import it lazily.
+    if name == "RedisStore":
+        from .redis import RedisStore
+
+        return RedisStore
+    raise AttributeError(name)
+
+
+__all__ = [
+    "LockTimeout",
+    "MemoryStore",
+    "MeterTable",
+    "RateLimitDecision",
+    "RedisStore",
+    "SlidingWindowCounter",
+    "StateStore",
+    "TokenBucket",
+]

wap/storage/base.py

@@ -0,0 +1,70 @@
+"""The storage interface every piece of server state goes through.
+
+A WAP server keeps several kinds of short-lived state: rate-limit meters,
+replay caches, spent proof-of-work seeds, dialogue sessions, loop-guard
+history and idempotency records. By default all of it lives in process memory
+(:class:`~wap.storage.memory.MemoryStore`), which is right for a single worker.
+Deployments with several workers or machines pass a shared backend such as
+:class:`~wap.storage.redis.RedisStore` so that limits, replay protection and
+sessions hold across all of them::
+
+    from wap.storage import RedisStore
+    wap = WAPServer(..., store=RedisStore.from_url("redis://localhost:6379/0"))
+
+Values are JSON documents. Implementations must make :meth:`StateStore.add`
+and :meth:`StateStore.rate_limit` atomic across every process sharing the store.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Protocol, runtime_checkable
+
+from .meters import RateLimitDecision
+
+
+class LockTimeout(Exception):
+    """A named lock could not be acquired within the allowed wait."""
+
+
+@runtime_checkable
+class StateStore(Protocol):
+    """Async key/value store with TTLs, atomic set-if-absent, locks and rate limiting."""
+
+    async def get(self, key: str) -> Any | None:
+        """Return the JSON value stored at ``key`` (``None`` if missing or expired)."""
+        ...
+
+    async def set(self, key: str, value: Any, ttl: float | None = None) -> None:
+        """Store a JSON-serialisable ``value``, expiring after ``ttl`` seconds if given."""
+        ...
+
+    async def add(self, key: str, value: Any, ttl: float | None = None) -> bool:
+        """Atomically store ``value`` only if ``key`` is absent. Returns ``True`` if stored."""
+        ...
+
+    async def delete(self, key: str) -> None: ...
+
+    async def incr(self, key: str, ttl: float) -> int:
+        """Atomically increment an integer counter, creating it with ``ttl`` if absent."""
+        ...
+
+    def lock(self, key: str, *, timeout: float, wait: float) -> AbstractAsyncContextManager[None]:
+        """Mutual exclusion on ``key`` across all processes sharing the store.
+
+        ``timeout`` bounds how long the lock may be held (so a crashed holder
+        cannot block forever); ``wait`` bounds how long to wait for it before
+        raising :class:`LockTimeout`.
+        """
+        ...
+
+    async def rate_limit(
+        self, keys: list[tuple[str, str]], *, limit: int, window: float, burst: int, cost: float, now: float
+    ) -> RateLimitDecision:
+        """Atomically check and (if admitted) charge every ``(scope, key)`` meter."""
+        ...
+
+    async def aclose(self) -> None: ...
+
+
+__all__ = ["LockTimeout", "StateStore"]

wap/storage/memory.py

@@ -0,0 +1,115 @@
+"""In-process :class:`StateStore`: the zero-configuration default for single-worker servers."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from collections import OrderedDict
+from collections.abc import AsyncIterator, Callable
+from contextlib import asynccontextmanager
+from typing import Any
+
+from .base import LockTimeout
+from .meters import MeterTable, RateLimitDecision
+
+
+class MemoryStore:
+    """Dictionary-backed store with TTLs, LRU bounding and asyncio locks.
+
+    Values are stored as JSON text, exactly as a networked backend would, so code
+    that works against this store behaves the same against Redis. State is lost on
+    restart and is not shared between processes; use a shared backend for that.
+    """
+
+    def __init__(
+        self,
+        *,
+        max_entries: int = 200_000,
+        max_rate_keys: int = 100_000,
+        rate_idle_ttl: float = 600.0,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        self.max_entries = max_entries
+        self._clock = clock
+        self._data: OrderedDict[str, tuple[str, float | None]] = OrderedDict()
+        self._locks: dict[str, asyncio.Lock] = {}
+        self.meters = MeterTable(idle_ttl=rate_idle_ttl, max_keys=max_rate_keys)
+        self._meter_lock = asyncio.Lock()
+
+    def _live(self, key: str) -> str | None:
+        entry = self._data.get(key)
+        if entry is None:
+            return None
+        raw, expires = entry
+        if expires is not None and expires <= self._clock():
+            del self._data[key]
+            return None
+        return raw
+
+    def _put(self, key: str, value: Any, ttl: float | None) -> None:
+        expires = self._clock() + ttl if ttl is not None else None
+        self._data[key] = (json.dumps(value, separators=(",", ":")), expires)
+        self._data.move_to_end(key)
+        if len(self._data) > self.max_entries:
+            now = self._clock()
+            for stale in [k for k, (_, exp) in self._data.items() if exp is not None and exp <= now]:
+                del self._data[stale]
+            while len(self._data) > self.max_entries:
+                self._data.popitem(last=False)
+
+    async def get(self, key: str) -> Any | None:
+        raw = self._live(key)
+        return None if raw is None else json.loads(raw)
+
+    async def set(self, key: str, value: Any, ttl: float | None = None) -> None:
+        self._put(key, value, ttl)
+
+    async def add(self, key: str, value: Any, ttl: float | None = None) -> bool:
+        if self._live(key) is not None:
+            return False
+        self._put(key, value, ttl)
+        return True
+
+    async def delete(self, key: str) -> None:
+        self._data.pop(key, None)
+
+    async def incr(self, key: str, ttl: float) -> int:
+        raw = self._live(key)
+        if raw is None:
+            self._put(key, 1, ttl)
+            return 1
+        value = int(json.loads(raw)) + 1
+        expires = self._data[key][1]
+        self._data[key] = (json.dumps(value), expires)
+        return value
+
+    @asynccontextmanager
+    async def lock(self, key: str, *, timeout: float, wait: float) -> AsyncIterator[None]:
+        lock = self._locks.setdefault(key, asyncio.Lock())
+        try:
+            await asyncio.wait_for(lock.acquire(), timeout=wait)
+        except TimeoutError as exc:
+            raise LockTimeout(key) from exc
+        try:
+            yield
+        finally:
+            lock.release()
+            if not lock.locked() and not getattr(lock, "_waiters", None):
+                self._locks.pop(key, None)
+
+    async def rate_limit(
+        self, keys: list[tuple[str, str]], *, limit: int, window: float, burst: int, cost: float, now: float
+    ) -> RateLimitDecision:
+        async with self._meter_lock:
+            return self.meters.check(keys, limit=limit, window=window, burst=burst, cost=cost, now=now)
+
+    async def aclose(self) -> None:
+        self._data.clear()
+        self.meters.clear()
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+
+__all__ = ["MemoryStore"]

wap/storage/meters.py

@@ -0,0 +1,204 @@
+"""Rate-limiting algorithms shared by the storage backends.
+
+Each principal gets two meters:
+
+* a **sliding-window counter** bounding the sustained rate (``limit`` requests
+  over a rolling ``window``), and
+* a **token bucket** bounding bursts (``burst`` tokens refilled at
+  ``limit / window`` per second).
+
+A request is admitted only if *every* meter it touches admits it, and meters
+are charged only on admission. :class:`MeterTable` is the in-process
+implementation; the Redis backend runs the same arithmetic in a Lua script.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class RateLimitDecision:
+    allowed: bool
+    limit: int
+    remaining: int
+    retry_after: float
+    scope: str | None = None
+
+    def headers(self) -> dict[str, str]:
+        """Standard ``RateLimit-*`` style headers for this decision."""
+        headers = {
+            "X-RateLimit-Limit": str(self.limit),
+            "X-RateLimit-Remaining": str(max(0, self.remaining)),
+        }
+        if not self.allowed:
+            headers["Retry-After"] = str(max(1, math.ceil(self.retry_after)))
+        return headers
+
+
+class TokenBucket:
+    """Classic token bucket: ``capacity`` tokens, refilled at ``rate`` tokens/second."""
+
+    __slots__ = ("capacity", "rate", "tokens", "updated")
+
+    def __init__(self, capacity: float, rate: float, now: float) -> None:
+        self.capacity = float(capacity)
+        self.rate = float(rate)
+        self.tokens = float(capacity)
+        self.updated = now
+
+    def _refill(self, now: float) -> None:
+        elapsed = max(0.0, now - self.updated)
+        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
+        self.updated = now
+
+    def peek(self, now: float, cost: float = 1.0) -> float:
+        """Seconds until ``cost`` tokens are available (0 if available now)."""
+        self._refill(now)
+        if self.tokens >= cost:
+            return 0.0
+        return (cost - self.tokens) / self.rate
+
+    def consume(self, now: float, cost: float = 1.0) -> None:
+        self._refill(now)
+        self.tokens -= cost
+
+    @property
+    def remaining(self) -> int:
+        return int(self.tokens)
+
+
+class SlidingWindowCounter:
+    """Sliding-window counter (Cloudflare-style weighted two-bucket approximation).
+
+    The estimated count over the last ``window`` seconds is
+    ``previous * (1 - elapsed_fraction) + current``. Memory is O(1) per key and
+    the error versus a true sliding log is bounded by one window's skew.
+    """
+
+    __slots__ = ("limit", "window", "window_start", "current", "previous")
+
+    def __init__(self, limit: int, window: float, now: float) -> None:
+        self.limit = limit
+        self.window = window
+        self.window_start = now - (now % window)
+        self.current = 0
+        self.previous = 0
+
+    def _roll(self, now: float) -> None:
+        start = now - (now % self.window)
+        if start == self.window_start:
+            return
+        if start - self.window_start >= 2 * self.window:
+            self.previous = 0
+        else:
+            self.previous = self.current
+        self.current = 0
+        self.window_start = start
+
+    def estimate(self, now: float) -> float:
+        self._roll(now)
+        elapsed = (now - self.window_start) / self.window
+        return self.previous * (1.0 - elapsed) + self.current
+
+    def peek(self, now: float) -> float:
+        """Seconds until one more request fits (0 if it fits now)."""
+        estimate = self.estimate(now)
+        if estimate + 1 <= self.limit:
+            return 0.0
+        if self.current + 1 > self.limit:
+            return self.window_start + self.window - now
+        if self.previous == 0:
+            return self.window_start + self.window - now
+        # Solve previous * (1 - (t - start)/window) + current + 1 <= limit for t.
+        needed_fraction = 1.0 - (self.limit - self.current - 1) / self.previous
+        target = self.window_start + needed_fraction * self.window
+        return max(0.0, target - now)
+
+    def consume(self, now: float) -> None:
+        self._roll(now)
+        self.current += 1
+
+    def remaining(self, now: float) -> int:
+        return max(0, int(self.limit - self.estimate(now)))
+
+
+class Meter:
+    __slots__ = ("window", "bucket", "last_seen", "params")
+
+    def __init__(self, limit: int, window: float, burst: int, now: float) -> None:
+        self.window = SlidingWindowCounter(limit, window, now)
+        self.bucket = TokenBucket(burst, limit / window, now)
+        self.last_seen = now
+        self.params = (limit, window, burst)
+
+    def wait_time(self, now: float, cost: float) -> float:
+        return max(self.window.peek(now), self.bucket.peek(now, cost))
+
+    def consume(self, now: float, cost: float) -> None:
+        self.window.consume(now)
+        self.bucket.consume(now, cost)
+        self.last_seen = now
+
+    def remaining(self, now: float) -> int:
+        return min(self.window.remaining(now), self.bucket.remaining)
+
+
+class MeterTable:
+    """Bounded in-process table of meters with idle eviction."""
+
+    def __init__(self, *, idle_ttl: float = 600.0, max_keys: int = 100_000) -> None:
+        self.idle_ttl = idle_ttl
+        self.max_keys = max_keys
+        self._meters: dict[str, Meter] = {}
+        self._last_gc: float | None = None
+
+    def _gc(self, now: float, window: float) -> None:
+        if self._last_gc is None:
+            self._last_gc = now
+        if now - self._last_gc < window and len(self._meters) < self.max_keys:
+            return
+        cutoff = now - self.idle_ttl
+        for key in [k for k, m in self._meters.items() if m.last_seen < cutoff]:
+            del self._meters[key]
+        if len(self._meters) >= self.max_keys:
+            # Evict least-recently-seen principals to bound memory under key-spraying.
+            ordered = sorted(self._meters.items(), key=lambda item: item[1].last_seen)
+            for key, _ in ordered[: len(self._meters) - self.max_keys + 1]:
+                del self._meters[key]
+        self._last_gc = now
+
+    def check(
+        self, keys: list[tuple[str, str]], *, limit: int, window: float, burst: int, cost: float, now: float
+    ) -> RateLimitDecision:
+        self._gc(now, window)
+        meters = []
+        for scope, key in keys:
+            meter = self._meters.get(key)
+            if meter is None or meter.params != (limit, window, burst):
+                meter = self._meters[key] = Meter(limit, window, burst, now)
+            meters.append((scope, meter))
+        worst_scope, worst_wait = None, 0.0
+        for scope, meter in meters:
+            wait = meter.wait_time(now, cost)
+            if wait > worst_wait:
+                worst_scope, worst_wait = scope, wait
+        if worst_wait > 0:
+            remaining = min(m.remaining(now) for _, m in meters)
+            return RateLimitDecision(
+                allowed=False, limit=limit, remaining=remaining, retry_after=worst_wait, scope=worst_scope
+            )
+        for _, meter in meters:
+            meter.consume(now, cost)
+        remaining = min(m.remaining(now) for _, m in meters)
+        return RateLimitDecision(allowed=True, limit=limit, remaining=remaining, retry_after=0.0)
+
+    def clear(self) -> None:
+        self._meters.clear()
+
+    def __len__(self) -> int:
+        return len(self._meters)
+
+
+__all__ = ["Meter", "MeterTable", "RateLimitDecision", "SlidingWindowCounter", "TokenBucket"]

wap/storage/redis.py

@@ -0,0 +1,202 @@
+"""Redis-backed :class:`StateStore` for multi-worker and multi-machine deployments.
+
+Install with ``pip install "webagent-protocol[redis]"`` and pass it to the server::
+
+    from wap.storage import RedisStore
+    wap = WAPServer(..., store=RedisStore.from_url("redis://redis:6379/0", prefix="bakery:"))
+
+All workers pointing at the same Redis database (and prefix) share rate limits,
+replay protection, spent proof-of-work seeds, sessions, loop guards and
+idempotency records. Rate limiting runs as a single Lua script, so the
+check-and-charge of every meter involved in a request is atomic.
+
+Keys are namespaced with ``prefix``. Multi-key rate-limit calls must hash to one
+slot on Redis Cluster; use a single-shard deployment or a ``prefix`` containing a
+hash tag such as ``"{wap}:"`` if you run a cluster.
+"""
+
+from __future__ import annotations
+
+import json
+import secrets
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+try:
+    from redis.asyncio import Redis
+except ImportError as exc:  # pragma: no cover - exercised only without the extra
+    raise ImportError('RedisStore requires the redis extra: pip install "webagent-protocol[redis]"') from exc
+
+from .base import LockTimeout
+from .meters import RateLimitDecision
+
+# KEYS: meter keys. ARGV: limit, window, burst, cost, now, idle_ttl_ms.
+# Returns {allowed(0/1), retry_after(string), worst_index, remaining}.
+_RATE_LIMIT_LUA = """
+local limit = tonumber(ARGV[1])
+local window = tonumber(ARGV[2])
+local burst = tonumber(ARGV[3])
+local cost = tonumber(ARGV[4])
+local now = tonumber(ARGV[5])
+local idle_ms = tonumber(ARGV[6])
+local rate = limit / window
+local state = {}
+local worst, worst_i = 0, 0
+for i, key in ipairs(KEYS) do
+  local v = redis.call('HMGET', key, 'ws', 'cur', 'prev', 'tokens', 'upd')
+  local ws, cur, prev, tokens, upd
+  if v[1] then
+    ws, cur, prev = tonumber(v[1]), tonumber(v[2]), tonumber(v[3])
+    tokens, upd = tonumber(v[4]), tonumber(v[5])
+  else
+    ws, cur, prev, tokens, upd = now - (now % window), 0, 0, burst, now
+  end
+  local start = now - (now % window)
+  if start ~= ws then
+    if start - ws >= 2 * window then prev = 0 else prev = cur end
+    cur = 0
+    ws = start
+  end
+  local elapsed = now - upd
+  if elapsed < 0 then elapsed = 0 end
+  tokens = math.min(burst, tokens + elapsed * rate)
+  upd = now
+  local estimate = prev * (1 - (now - ws) / window) + cur
+  local wait_w = 0
+  if estimate + 1 > limit then
+    if cur + 1 > limit or prev == 0 then
+      wait_w = ws + window - now
+    else
+      local needed = 1 - (limit - cur - 1) / prev
+      wait_w = math.max(0, ws + needed * window - now)
+    end
+  end
+  local wait_b = 0
+  if tokens < cost then wait_b = (cost - tokens) / rate end
+  local wait = math.max(wait_w, wait_b)
+  if wait > worst then worst, worst_i = wait, i end
+  state[i] = {ws, cur, prev, tokens, upd}
+end
+local allowed = 1
+if worst > 0 then allowed = 0 end
+local remaining = nil
+for i, key in ipairs(KEYS) do
+  local s = state[i]
+  if allowed == 1 then
+    s[2] = s[2] + 1
+    s[4] = s[4] - cost
+  end
+  redis.call('HSET', key, 'ws', tostring(s[1]), 'cur', tostring(s[2]), 'prev', tostring(s[3]),
+             'tokens', tostring(s[4]), 'upd', tostring(s[5]))
+  redis.call('PEXPIRE', key, idle_ms)
+  local est = s[3] * (1 - (now - s[1]) / window) + s[2]
+  local r = math.min(math.floor(limit - est), math.floor(s[4]))
+  if r < 0 then r = 0 end
+  if remaining == nil or r < remaining then remaining = r end
+end
+return {allowed, tostring(worst), worst_i, remaining}
+"""
+
+_RELEASE_LUA = """
+if redis.call('GET', KEYS[1]) == ARGV[1] then return redis.call('DEL', KEYS[1]) end
+return 0
+"""
+
+
+class RedisStore:
+    """Shared state in Redis. Safe to use from many workers and machines at once."""
+
+    def __init__(
+        self,
+        client: Redis,
+        *,
+        prefix: str = "wap:",
+        rate_idle_ttl: float = 600.0,
+        lock_poll_interval: float = 0.05,
+    ) -> None:
+        self.redis = client
+        self.prefix = prefix
+        self.rate_idle_ttl = rate_idle_ttl
+        self.lock_poll_interval = lock_poll_interval
+        self._rate_script = client.register_script(_RATE_LIMIT_LUA)
+        self._release_script = client.register_script(_RELEASE_LUA)
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs: Any) -> RedisStore:
+        """Connect with a ``redis://`` / ``rediss://`` URL (``rediss`` for TLS)."""
+        store_kwargs = {k: kwargs.pop(k) for k in ("prefix", "rate_idle_ttl", "lock_poll_interval") if k in kwargs}
+        return cls(Redis.from_url(url, **kwargs), **store_kwargs)
+
+    def _k(self, key: str) -> str:
+        return self.prefix + key
+
+    async def get(self, key: str) -> Any | None:
+        raw = await self.redis.get(self._k(key))
+        return None if raw is None else json.loads(raw)
+
+    async def set(self, key: str, value: Any, ttl: float | None = None) -> None:
+        px = max(1, int(ttl * 1000)) if ttl is not None else None
+        await self.redis.set(self._k(key), json.dumps(value, separators=(",", ":")), px=px)
+
+    async def add(self, key: str, value: Any, ttl: float | None = None) -> bool:
+        px = max(1, int(ttl * 1000)) if ttl is not None else None
+        stored = await self.redis.set(self._k(key), json.dumps(value, separators=(",", ":")), px=px, nx=True)
+        return bool(stored)
+
+    async def delete(self, key: str) -> None:
+        await self.redis.delete(self._k(key))
+
+    async def incr(self, key: str, ttl: float) -> int:
+        name = self._k(key)
+        pipe = self.redis.pipeline(transaction=True)
+        pipe.incr(name)
+        pipe.pexpire(name, max(1, int(ttl * 1000)), nx=True)
+        value, _ = await pipe.execute()
+        return int(value)
+
+    @asynccontextmanager
+    async def lock(self, key: str, *, timeout: float, wait: float) -> AsyncIterator[None]:
+        import anyio
+
+        name = self._k("lock:" + key)
+        token = secrets.token_hex(16)
+        px = max(1, int(timeout * 1000))
+        acquired = False
+        with anyio.move_on_after(wait):
+            while not acquired:
+                acquired = bool(await self.redis.set(name, token, px=px, nx=True))
+                if not acquired:
+                    await anyio.sleep(self.lock_poll_interval)
+        if not acquired:
+            raise LockTimeout(key)
+        try:
+            yield
+        finally:
+            # Only release if we still own it (the lock may have expired and been re-taken).
+            await self._release_script(keys=[name], args=[token])
+
+    async def rate_limit(
+        self, keys: list[tuple[str, str]], *, limit: int, window: float, burst: int, cost: float, now: float
+    ) -> RateLimitDecision:
+        if not keys:
+            return RateLimitDecision(allowed=True, limit=limit, remaining=limit, retry_after=0.0)
+        redis_keys = [self._k(f"rl:{limit}:{window}:{burst}:{key}") for _, key in keys]
+        allowed, retry_after, worst_index, remaining = await self._rate_script(
+            keys=redis_keys,
+            args=[limit, window, burst, cost, repr(float(now)), int(self.rate_idle_ttl * 1000)],
+        )
+        scope = keys[int(worst_index) - 1][0] if int(worst_index) > 0 else None
+        return RateLimitDecision(
+            allowed=bool(int(allowed)),
+            limit=limit,
+            remaining=int(remaining),
+            retry_after=float(retry_after),
+            scope=None if int(allowed) else scope,
+        )
+
+    async def aclose(self) -> None:
+        await self.redis.aclose()
+
+
+__all__ = ["RedisStore"]