tollgate 1.0.4__py3-none-any.whl → 1.4.0__py3-none-any.whl

tollgate/backends/redis_store.py

@@ -0,0 +1,411 @@
+"""Redis-backed persistent stores for Tollgate.
+
+Requires the ``redis`` package: ``pip install redis[hiredis]``
+
+Suitable for multi-process and multi-host deployments. Uses Redis hashes
+for grant storage and pub/sub for real-time approval notifications.
+
+Usage:
+
+    from tollgate.backends import RedisGrantStore, RedisApprovalStore
+
+    grant_store = RedisGrantStore(redis_url="redis://localhost:6379/0")
+    approval_store = RedisApprovalStore(redis_url="redis://localhost:6379/0")
+
+    tower = ControlTower(
+        ...,
+        grant_store=grant_store,
+    )
+"""
+
+import asyncio
+import json
+import time
+import uuid
+from typing import Any
+
+try:
+    import redis.asyncio as aioredis
+except ImportError:
+    raise ImportError(
+        "Redis backend requires the 'redis' package. "
+        "Install it with: pip install redis[hiredis]"
+    )
+
+from ..types import AgentContext, ApprovalOutcome, Effect, Grant, ToolRequest
+
+
+class RedisGrantStore:
+    """Redis-backed GrantStore implementation.
+
+    Uses Redis hashes for grant storage with automatic TTL-based expiry.
+    Satisfies the ``GrantStore`` protocol.
+
+    Args:
+        redis_url: Redis connection URL (e.g., ``redis://localhost:6379/0``).
+        redis_client: Pre-configured async Redis client (alternative to URL).
+        key_prefix: Prefix for all Redis keys (default ``tollgate:grant:``).
+    """
+
+    def __init__(
+        self,
+        redis_url: str | None = None,
+        *,
+        redis_client: Any | None = None,
+        key_prefix: str = "tollgate:grant:",
+    ):
+        if redis_client is not None:
+            self._redis = redis_client
+        elif redis_url is not None:
+            self._redis = aioredis.from_url(redis_url, decode_responses=True)
+        else:
+            raise ValueError("Either redis_url or redis_client must be provided")
+
+        self._prefix = key_prefix
+        self._index_key = f"{key_prefix}__index__"
+
+    def _grant_key(self, grant_id: str) -> str:
+        return f"{self._prefix}{grant_id}"
+
+    def _grant_to_dict(self, grant: Grant) -> dict[str, str]:
+        """Serialize a Grant to a flat dict for Redis HSET."""
+        return {
+            "id": grant.id,
+            "agent_id": grant.agent_id or "",
+            "effect": grant.effect.value if grant.effect else "",
+            "tool": grant.tool or "",
+            "action": grant.action or "",
+            "resource_type": grant.resource_type or "",
+            "expires_at": str(grant.expires_at),
+            "granted_by": grant.granted_by,
+            "created_at": str(grant.created_at),
+            "reason": grant.reason or "",
+            "usage_count": "0",
+            "revoked": "0",
+        }
+
+    def _dict_to_grant(self, d: dict[str, str]) -> Grant:
+        """Deserialize a Redis hash dict to a Grant."""
+        return Grant(
+            id=d["id"],
+            agent_id=d["agent_id"] or None,
+            effect=Effect(d["effect"]) if d["effect"] else None,
+            tool=d["tool"] or None,
+            action=d["action"] or None,
+            resource_type=d["resource_type"] or None,
+            expires_at=float(d["expires_at"]),
+            granted_by=d["granted_by"],
+            created_at=float(d["created_at"]),
+            reason=d["reason"] or None,
+        )
+
+    async def create_grant(self, grant: Grant) -> str:
+        key = self._grant_key(grant.id)
+        data = self._grant_to_dict(grant)
+
+        pipe = self._redis.pipeline()
+        pipe.hset(key, mapping=data)
+        # Set TTL based on expiry time
+        ttl = max(1, int(grant.expires_at - time.time()))
+        pipe.expire(key, ttl)
+        # Track in index set
+        pipe.sadd(self._index_key, grant.id)
+        await pipe.execute()
+
+        return grant.id
+
+    async def find_matching_grant(
+        self, agent_ctx: AgentContext, tool_request: ToolRequest
+    ) -> Grant | None:
+        # Get all active grant IDs from the index
+        grant_ids = await self._redis.smembers(self._index_key)
+
+        for grant_id in grant_ids:
+            key = self._grant_key(grant_id)
+            data = await self._redis.hgetall(key)
+
+            if not data:
+                # Key expired — remove from index
+                await self._redis.srem(self._index_key, grant_id)
+                continue
+
+            if data.get("revoked") == "1":
+                continue
+
+            expires_at = float(data["expires_at"])
+            if expires_at <= time.time():
+                continue
+
+            # Match agent_id
+            if data["agent_id"] and data["agent_id"] != agent_ctx.agent_id:
+                continue
+
+            # Match effect
+            if data["effect"] and data["effect"] != tool_request.effect.value:
+                continue
+
+            # Match tool (exact or prefix with *)
+            if data["tool"]:
+                if data["tool"].endswith("*"):
+                    prefix = data["tool"][:-1]
+                    if not tool_request.tool.startswith(prefix):
+                        continue
+                elif data["tool"] != tool_request.tool:
+                    continue
+
+            # Match action
+            if data["action"] and data["action"] != tool_request.action:
+                continue
+
+            # Match resource_type
+            if (
+                data["resource_type"]
+                and data["resource_type"] != tool_request.resource_type
+            ):
+                continue
+
+            # Match found — increment usage
+            await self._redis.hincrby(key, "usage_count", 1)
+            return self._dict_to_grant(data)
+
+        return None
+
+    async def revoke_grant(self, grant_id: str) -> bool:
+        key = self._grant_key(grant_id)
+        exists = await self._redis.exists(key)
+        if not exists:
+            return False
+
+        await self._redis.hset(key, "revoked", "1")
+        return True
+
+    async def list_active_grants(self, agent_id: str | None = None) -> list[Grant]:
+        grant_ids = await self._redis.smembers(self._index_key)
+        now = time.time()
+        active: list[Grant] = []
+
+        for grant_id in grant_ids:
+            key = self._grant_key(grant_id)
+            data = await self._redis.hgetall(key)
+
+            if not data:
+                await self._redis.srem(self._index_key, grant_id)
+                continue
+
+            if data.get("revoked") == "1":
+                continue
+
+            if float(data["expires_at"]) <= now:
+                continue
+
+            if agent_id is not None and data["agent_id"] != agent_id:
+                continue
+
+            active.append(self._dict_to_grant(data))
+
+        return active
+
+    async def cleanup_expired(self) -> int:
+        grant_ids = await self._redis.smembers(self._index_key)
+        now = time.time()
+        removed = 0
+
+        for grant_id in grant_ids:
+            key = self._grant_key(grant_id)
+            data = await self._redis.hgetall(key)
+
+            if not data or float(data.get("expires_at", "0")) <= now:
+                await self._redis.delete(key)
+                await self._redis.srem(self._index_key, grant_id)
+                removed += 1
+
+        return removed
+
+    async def get_usage_count(self, grant_id: str) -> int:
+        key = self._grant_key(grant_id)
+        count = await self._redis.hget(key, "usage_count")
+        return int(count) if count else 0
+
+    async def close(self):
+        """Close the Redis connection."""
+        await self._redis.close()
+
+
+class RedisApprovalStore:
+    """Redis-backed ApprovalStore implementation.
+
+    Uses Redis hashes for request storage and pub/sub for real-time
+    notifications when a decision is made. This avoids polling.
+
+    Args:
+        redis_url: Redis connection URL.
+        redis_client: Pre-configured async Redis client.
+        key_prefix: Prefix for all Redis keys.
+        channel_prefix: Prefix for pub/sub channels.
+    """
+
+    def __init__(
+        self,
+        redis_url: str | None = None,
+        *,
+        redis_client: Any | None = None,
+        key_prefix: str = "tollgate:approval:",
+        channel_prefix: str = "tollgate:approval_notify:",
+    ):
+        if redis_client is not None:
+            self._redis = redis_client
+        elif redis_url is not None:
+            self._redis = aioredis.from_url(redis_url, decode_responses=True)
+        else:
+            raise ValueError("Either redis_url or redis_client must be provided")
+
+        self._prefix = key_prefix
+        self._channel_prefix = channel_prefix
+
+    def _request_key(self, approval_id: str) -> str:
+        return f"{self._prefix}{approval_id}"
+
+    def _channel_key(self, approval_id: str) -> str:
+        return f"{self._channel_prefix}{approval_id}"
+
+    async def create_request(
+        self,
+        agent_ctx: AgentContext,
+        intent: Any,
+        tool_request: ToolRequest,
+        request_hash: str,
+        reason: str,
+        expiry: float,
+    ) -> str:
+        approval_id = str(uuid.uuid4())
+        key = self._request_key(approval_id)
+
+        data = {
+            "id": approval_id,
+            "agent_json": json.dumps(agent_ctx.to_dict()),
+            "intent_json": json.dumps(intent.to_dict()),
+            "tool_request_json": json.dumps(tool_request.to_dict()),
+            "request_hash": request_hash,
+            "reason": reason,
+            "expiry": str(expiry),
+            "outcome": ApprovalOutcome.DEFERRED.value,
+            "decided_by": "",
+            "decided_at": "",
+        }
+
+        pipe = self._redis.pipeline()
+        pipe.hset(key, mapping=data)
+        # Auto-expire after the approval window
+        ttl = max(1, int(expiry - time.time()) + 60)  # +60s buffer
+        pipe.expire(key, ttl)
+        await pipe.execute()
+
+        return approval_id
+
+    async def set_decision(
+        self,
+        approval_id: str,
+        outcome: ApprovalOutcome,
+        decided_by: str,
+        decided_at: float,
+        request_hash: str,
+    ) -> None:
+        key = self._request_key(approval_id)
+
+        # Verify request hash (replay protection)
+        stored_hash = await self._redis.hget(key, "request_hash")
+        if stored_hash is None:
+            return  # Request not found
+
+        if stored_hash != request_hash:
+            raise ValueError(
+                "Request hash mismatch. Approval bound to a different request."
+            )
+
+        pipe = self._redis.pipeline()
+        pipe.hset(key, mapping={
+            "outcome": outcome.value,
+            "decided_by": decided_by,
+            "decided_at": str(decided_at),
+        })
+        # Publish notification for waiters
+        pipe.publish(self._channel_key(approval_id), outcome.value)
+        await pipe.execute()
+
+    async def get_request(self, approval_id: str) -> dict[str, Any] | None:
+        key = self._request_key(approval_id)
+        data = await self._redis.hgetall(key)
+
+        if not data:
+            return None
+
+        return {
+            "id": data["id"],
+            "agent": json.loads(data["agent_json"]),
+            "intent": json.loads(data["intent_json"]),
+            "tool_request": json.loads(data["tool_request_json"]),
+            "request_hash": data["request_hash"],
+            "reason": data["reason"],
+            "expiry": float(data["expiry"]),
+            "outcome": ApprovalOutcome(data["outcome"]),
+            "decided_by": data.get("decided_by") or None,
+            "decided_at": float(data["decided_at"]) if data.get("decided_at") else None,
+        }
+
+    async def wait_for_decision(
+        self, approval_id: str, timeout: float
+    ) -> ApprovalOutcome:
+        """Wait for a decision using Redis pub/sub (non-polling)."""
+        # First check if a decision already exists
+        req = await self.get_request(approval_id)
+        if req is None:
+            return ApprovalOutcome.TIMEOUT
+
+        if req["expiry"] < time.time():
+            return ApprovalOutcome.TIMEOUT
+
+        if req["outcome"] != ApprovalOutcome.DEFERRED:
+            return req["outcome"]
+
+        # Subscribe and wait for a notification
+        channel = self._channel_key(approval_id)
+        pubsub = self._redis.pubsub()
+
+        try:
+            await pubsub.subscribe(channel)
+
+            deadline = time.time() + timeout
+            while time.time() < deadline:
+                remaining = deadline - time.time()
+                if remaining <= 0:
+                    break
+
+                try:
+                    message = await asyncio.wait_for(
+                        pubsub.get_message(
+                            ignore_subscribe_messages=True, timeout=1.0
+                        ),
+                        timeout=min(remaining, 2.0),
+                    )
+                except asyncio.TimeoutError:
+                    # Check if decision was made (fallback polling)
+                    req = await self.get_request(approval_id)
+                    if req and req["outcome"] != ApprovalOutcome.DEFERRED:
+                        return req["outcome"]
+                    continue
+
+                if message and message["type"] == "message":
+                    try:
+                        return ApprovalOutcome(message["data"])
+                    except ValueError:
+                        pass
+
+            return ApprovalOutcome.TIMEOUT
+
+        finally:
+            await pubsub.unsubscribe(channel)
+            await pubsub.close()
+
+    async def close(self):
+        """Close the Redis connection."""
+        await self._redis.close()