"""Per-installation sliding-window rate limiting backed by Redis. Three configurable buckets, all opt-in via env vars: - ``DEEPAGENT_QUOTA_HOUR`` — max jobs accepted per installation per hour (default 0 = off) - ``DEEPAGENT_QUOTA_DAY`` — max jobs per installation per day (default 0 = off) - ``DEEPAGENT_QUOTA_CONCURRENT`` — max concurrent in-flight jobs per installation (default 0 = off) Per-installation overrides via the env var ``DEEPAGENT_QUOTA_OVERRIDES`` (JSON map of ``"": {"hour": N, "day": N, "concurrent": N}``). The webhook server consults :class:`QuotaManager.check_and_consume` BEFORE enqueueing a job. If denied, it returns HTTP 429 with the limiting bucket name + retry-after seconds. The worker is responsible for calling :meth:`release_concurrent` when a job finishes (regardless of outcome). Implementation: ZSET per (bucket, installation_id) with score=timestamp. Old entries are trimmed with ZREMRANGEBYSCORE before counting. Concurrent counter is a plain INCR/DECR. """ from __future__ import annotations import json import logging import os import time from dataclasses import dataclass from typing import Optional import redis log = logging.getLogger(__name__) @dataclass class QuotaDecision: allowed: bool bucket: Optional[str] = None # "hour" | "day" | "concurrent" if denied retry_after_seconds: int = 0 current: int = 0 limit: int = 0 class QuotaManager: KEY_PREFIX = "deepagent:quota" BUCKETS = {"hour": 3600, "day": 86400} def __init__(self, client: Optional[redis.Redis] = None): self._r = client or redis.Redis.from_url( os.getenv("DEEPAGENT_REDIS_URL", "redis://localhost:6379/0") ) self._defaults = { "hour": int(os.getenv("DEEPAGENT_QUOTA_HOUR", "0")), "day": int(os.getenv("DEEPAGENT_QUOTA_DAY", "0")), "concurrent": int(os.getenv("DEEPAGENT_QUOTA_CONCURRENT", "0")), } self._overrides = self._load_overrides() @staticmethod def _load_overrides() -> dict[str, dict[str, int]]: raw = os.getenv("DEEPAGENT_QUOTA_OVERRIDES", "") if not raw: return {} try: return json.loads(raw) except json.JSONDecodeError: log.warning("DEEPAGENT_QUOTA_OVERRIDES not valid JSON; ignoring.") return {} # ---- limits resolution def limits_for(self, installation_id: str | int | None) -> dict[str, int]: if installation_id is None: return dict(self._defaults) ov = self._overrides.get(str(installation_id), {}) return {k: int(ov.get(k, self._defaults[k])) for k in self._defaults} # ---- key helpers def _bucket_key(self, bucket: str, installation_id: str) -> str: return f"{self.KEY_PREFIX}:{bucket}:{installation_id}" def _conc_key(self, installation_id: str) -> str: return f"{self.KEY_PREFIX}:concurrent:{installation_id}" # ---- public API def check_and_consume(self, installation_id: str | int | None) -> QuotaDecision: """Atomically check every enabled bucket; consume on success.""" if installation_id is None or installation_id == "": return QuotaDecision(allowed=True) iid = str(installation_id) limits = self.limits_for(iid) # Concurrent first — cheapest, no scan. conc_limit = limits["concurrent"] if conc_limit > 0: current = int(self._r.get(self._conc_key(iid)) or 0) if current >= conc_limit: return QuotaDecision(False, "concurrent", retry_after_seconds=10, current=current, limit=conc_limit) # Sliding windows. now = time.time() for bucket, window in self.BUCKETS.items(): lim = limits[bucket] if lim <= 0: continue key = self._bucket_key(bucket, iid) cutoff = now - window pipe = self._r.pipeline() pipe.zremrangebyscore(key, 0, cutoff) pipe.zcard(key) _, count = pipe.execute() count = int(count) if count >= lim: # Compute when the oldest entry will fall out of the window. oldest = self._r.zrange(key, 0, 0, withscores=True) retry = window if oldest: _, ts = oldest[0] retry = max(1, int(window - (now - float(ts)))) return QuotaDecision(False, bucket, retry_after_seconds=retry, current=count, limit=lim) # All checks passed → consume. pipe = self._r.pipeline() for bucket, window in self.BUCKETS.items(): if limits[bucket] > 0: key = self._bucket_key(bucket, iid) pipe.zadd(key, {f"{now}:{installation_id}": now}) pipe.expire(key, window + 60) if conc_limit > 0: pipe.incr(self._conc_key(iid)) pipe.expire(self._conc_key(iid), 3600) # safety auto-reset pipe.execute() return QuotaDecision(True) def release_concurrent(self, installation_id: str | int | None) -> None: """Decrement the concurrent counter. Called by the worker on job completion.""" if installation_id is None or installation_id == "": return iid = str(installation_id) if self.limits_for(iid)["concurrent"] <= 0: return try: val = self._r.decr(self._conc_key(iid)) if int(val) < 0: self._r.set(self._conc_key(iid), 0) except redis.RedisError as e: # pragma: no cover log.warning("release_concurrent failed: %s", e) def usage(self, installation_id: str | int) -> dict[str, dict]: """Inspection helper: current usage vs. limits for an installation.""" iid = str(installation_id) limits = self.limits_for(iid) out: dict[str, dict] = {} now = time.time() for bucket, window in self.BUCKETS.items(): key = self._bucket_key(bucket, iid) self._r.zremrangebyscore(key, 0, now - window) out[bucket] = {"used": int(self._r.zcard(key)), "limit": limits[bucket]} out["concurrent"] = { "used": int(self._r.get(self._conc_key(iid)) or 0), "limit": limits["concurrent"], } return out