contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/api/tenant_quotas.py

@@ -0,0 +1,231 @@
+"""
+Tenant Quotas — per-tenant resource limits and usage tracking for SaaS.
+
+Tracks and enforces:
+  - Graphs per tenant (default: 50)
+  - Nodes per tenant (default: 100,000)
+  - Ingestion jobs per hour (default: 20)
+  - LLM calls per hour (default: 200)
+  - Storage MB per tenant (default: 500)
+  - API requests per minute (default: 60)
+
+Usage:
+    from contextsynapse.api.tenant_quotas import get_quota_manager
+
+    quotas = get_quota_manager()
+    quotas.check("tenant_123", "graphs")       # raises 429 if over limit
+    quotas.record("tenant_123", "llm_calls")   # increment usage counter
+    usage = quotas.get_usage("tenant_123")      # get current usage
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import time
+from typing import Any, Dict, Optional
+
+from fastapi import HTTPException
+
+logger = logging.getLogger(__name__)
+
+# Default limits (can be overridden per tenant in tenant.config)
+DEFAULT_QUOTAS = {
+    "graphs": 50,              # max graphs per tenant
+    "nodes": 100_000,          # max total nodes across all graphs
+    "ingestion_per_hour": 20,  # max ingestion jobs per hour
+    "llm_calls_per_hour": 200, # max LLM API calls per hour
+    "storage_mb": 500,         # max disk storage in MB
+    "api_rpm": 60,             # max API requests per minute
+}
+
+
+class TenantQuotaManager:
+    """Redis-backed quota tracking and enforcement."""
+
+    KEY_PREFIX = "quota:"
+
+    def __init__(self, redis_url: str = None):
+        import redis
+        self._url = redis_url or os.environ.get("CONTEXTSYNAPSE_REDIS_URL") or os.environ.get("AICONTEXTDB_REDIS_URL", "redis://localhost:6379/0")
+        self._r = redis.from_url(self._url, decode_responses=True)
+
+    def _quota_key(self, tenant_id: str, resource: str) -> str:
+        return f"{self.KEY_PREFIX}{tenant_id}:{resource}"
+
+    def get_limits(self, tenant_id: str) -> Dict[str, int]:
+        """Get effective limits for a tenant (custom overrides + defaults)."""
+        # Check for custom limits in Redis
+        custom_key = f"{self.KEY_PREFIX}{tenant_id}:limits"
+        try:
+            custom = self._r.hgetall(custom_key)
+            if custom:
+                limits = dict(DEFAULT_QUOTAS)
+                for k, v in custom.items():
+                    try:
+                        limits[k] = int(v)
+                    except ValueError:
+                        pass
+                return limits
+        except Exception:
+            pass
+        return dict(DEFAULT_QUOTAS)
+
+    def set_limits(self, tenant_id: str, limits: Dict[str, int]):
+        """Set custom limits for a tenant (admin operation)."""
+        custom_key = f"{self.KEY_PREFIX}{tenant_id}:limits"
+        self._r.hset(custom_key, mapping={k: str(v) for k, v in limits.items()})
+
+    def check(self, tenant_id: str, resource: str, count: int = 1) -> bool:
+        """Check if tenant can use more of a resource. Raises 429 if over limit."""
+        limits = self.get_limits(tenant_id)
+        limit = limits.get(resource)
+        if limit is None:
+            return True  # no limit defined for this resource
+
+        current = self.get_current(tenant_id, resource)
+        if current + count > limit:
+            raise HTTPException(
+                status_code=429,
+                detail=f"Quota exceeded: {resource} ({current}/{limit}). Upgrade your plan for higher limits.",
+            )
+        return True
+
+    def get_current(self, tenant_id: str, resource: str) -> int:
+        """Get current usage count for a resource."""
+        key = self._quota_key(tenant_id, resource)
+        try:
+            # Time-windowed resources (per hour / per minute)
+            if resource.endswith("_per_hour"):
+                now = time.time()
+                self._r.zremrangebyscore(key, 0, now - 3600)
+                return self._r.zcard(key)
+            elif resource == "api_rpm":
+                now = time.time()
+                self._r.zremrangebyscore(key, 0, now - 60)
+                return self._r.zcard(key)
+            else:
+                # Absolute counters (graphs, nodes, storage)
+                val = self._r.get(key)
+                return int(val) if val else 0
+        except Exception:
+            return 0
+
+    def record(self, tenant_id: str, resource: str, count: int = 1):
+        """Record resource usage."""
+        key = self._quota_key(tenant_id, resource)
+        try:
+            if resource.endswith("_per_hour") or resource == "api_rpm":
+                # Time-windowed: add to sorted set with timestamp
+                import uuid
+                now = time.time()
+                pipe = self._r.pipeline()
+                for _ in range(count):
+                    pipe.zadd(key, {str(uuid.uuid4())[:8]: now})
+                ttl = 3600 if resource.endswith("_per_hour") else 120
+                pipe.expire(key, ttl)
+                pipe.execute()
+            else:
+                # Absolute: increment counter
+                self._r.incrby(key, count)
+        except Exception as e:
+            logger.debug("[QUOTA] Failed to record %s for %s: %s", resource, tenant_id, e)
+
+    def set_absolute(self, tenant_id: str, resource: str, value: int):
+        """Set an absolute usage value (e.g., after recounting nodes)."""
+        key = self._quota_key(tenant_id, resource)
+        try:
+            self._r.set(key, str(value))
+        except Exception:
+            pass
+
+    def get_usage(self, tenant_id: str) -> Dict[str, Any]:
+        """Get full usage report for a tenant."""
+        limits = self.get_limits(tenant_id)
+        usage = {}
+        for resource, limit in limits.items():
+            current = self.get_current(tenant_id, resource)
+            usage[resource] = {
+                "current": current,
+                "limit": limit,
+                "percent": round(current / limit * 100, 1) if limit > 0 else 0,
+                "remaining": max(0, limit - current),
+            }
+        return usage
+
+    def reset(self, tenant_id: str, resource: str = None):
+        """Reset usage counters (admin operation)."""
+        if resource:
+            key = self._quota_key(tenant_id, resource)
+            self._r.delete(key)
+        else:
+            # Reset all
+            pattern = f"{self.KEY_PREFIX}{tenant_id}:*"
+            cursor = 0
+            while True:
+                cursor, keys = self._r.scan(cursor, match=pattern, count=100)
+                keys = [k for k in keys if not k.endswith(":limits")]  # keep custom limits
+                if keys:
+                    self._r.delete(*keys)
+                if cursor == 0:
+                    break
+
+
+class InProcessQuotaManager:
+    """Fallback quota manager when Redis is unavailable (single-worker only)."""
+
+    def __init__(self):
+        self._counters: Dict[str, int] = {}
+
+    def get_limits(self, tenant_id: str) -> Dict[str, int]:
+        return dict(DEFAULT_QUOTAS)
+
+    def set_limits(self, tenant_id: str, limits: Dict[str, int]):
+        pass
+
+    def check(self, tenant_id: str, resource: str, count: int = 1) -> bool:
+        return True  # no enforcement without Redis
+
+    def get_current(self, tenant_id: str, resource: str) -> int:
+        return self._counters.get(f"{tenant_id}:{resource}", 0)
+
+    def record(self, tenant_id: str, resource: str, count: int = 1):
+        key = f"{tenant_id}:{resource}"
+        self._counters[key] = self._counters.get(key, 0) + count
+
+    def set_absolute(self, tenant_id: str, resource: str, value: int):
+        self._counters[f"{tenant_id}:{resource}"] = value
+
+    def get_usage(self, tenant_id: str) -> Dict[str, Any]:
+        limits = self.get_limits(tenant_id)
+        return {r: {"current": 0, "limit": l, "percent": 0, "remaining": l} for r, l in limits.items()}
+
+    def reset(self, tenant_id: str, resource: str = None):
+        pass
+
+
+# Singleton
+_manager = None
+
+
+def get_quota_manager():
+    """Get or create quota manager (Redis if available, in-process fallback)."""
+    global _manager
+    if _manager is not None:
+        return _manager
+
+    url = os.environ.get("CONTEXTSYNAPSE_REDIS_URL") or os.environ.get("AICONTEXTDB_REDIS_URL")
+    if url:
+        try:
+            mgr = TenantQuotaManager(url)
+            mgr._r.ping()
+            logger.info("Tenant quotas: Redis")
+            _manager = mgr
+            return _manager
+        except Exception:
+            pass
+
+    _manager = InProcessQuotaManager()
+    logger.info("Tenant quotas: in-process (no enforcement)")
+    return _manager

contextsynapse/api/tenants.py

@@ -0,0 +1,305 @@
+"""
+Tenant Registry
+===============
+Multi-tenant isolation layer.  Each tenant gets an API key and a set of
+namespaces that are invisible to other tenants.
+
+Usage::
+
+    registry = TenantRegistry()
+    tenant, api_key = registry.create("acme-corp")
+    # Caller uses: Authorization: Bearer <api_key>
+    # Internally, namespaces are stored as "{tenant_id}:{namespace}"
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import secrets
+import sqlite3
+from dataclasses import dataclass, field, asdict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TenantIdentity:
+    """A registered tenant (organisation / user)."""
+    tenant_id: str
+    name: str
+    status: str = "active"  # active, suspended, deleted
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    config: Dict[str, Any] = field(default_factory=dict)
+    # config may contain: max_graphs, max_sessions, rate_limit_rpm, etc.
+    api_key_hash: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        d = asdict(self)
+        d.pop("api_key_hash", None)
+        return d
+
+    def ns_prefix(self) -> str:
+        """Return the namespace prefix for this tenant."""
+        return f"{self.tenant_id}:"
+
+    def scoped_namespace(self, namespace: str) -> str:
+        """Convert a user-facing namespace to the internal tenant-scoped form."""
+        return f"{self.tenant_id}:{namespace}"
+
+    def unscoped_namespace(self, internal_ns: str) -> Optional[str]:
+        """Strip tenant prefix. Returns None if the namespace doesn't belong to this tenant."""
+        prefix = self.ns_prefix()
+        if internal_ns.startswith(prefix):
+            return internal_ns[len(prefix):]
+        return None
+
+
+class TenantRegistry:
+    """
+    Persistent tenant registry backed by SQLite.
+    Same pattern as AgentRegistry in context/agents.py.
+    """
+
+    _SCHEMA = """
+    CREATE TABLE IF NOT EXISTS tenants (
+        tenant_id TEXT PRIMARY KEY,
+        name TEXT NOT NULL UNIQUE,
+        status TEXT DEFAULT 'active',
+        created_at TEXT NOT NULL,
+        config TEXT DEFAULT '{}',
+        api_key_hash TEXT NOT NULL
+    );
+
+    CREATE INDEX IF NOT EXISTS idx_tenants_name ON tenants(name);
+    CREATE INDEX IF NOT EXISTS idx_tenants_status ON tenants(status);
+
+    CREATE TABLE IF NOT EXISTS api_keys (
+        key_id TEXT PRIMARY KEY,
+        tenant_id TEXT NOT NULL,
+        name TEXT NOT NULL,
+        key_hash TEXT NOT NULL,
+        scope TEXT DEFAULT 'read_write',
+        created_by TEXT,
+        created_at TEXT NOT NULL,
+        last_used_at TEXT,
+        expires_at TEXT,
+        status TEXT DEFAULT 'active'
+    );
+
+    CREATE INDEX IF NOT EXISTS idx_ak_tenant ON api_keys(tenant_id);
+    CREATE INDEX IF NOT EXISTS idx_ak_hash ON api_keys(key_hash);
+    """
+
+    def __init__(self, db_path: str = "contextcore_data/tenants.db"):
+        self._db_path = Path(db_path)
+        self._db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
+        self._conn.row_factory = sqlite3.Row
+        self._conn.executescript(self._SCHEMA)
+
+    # ------------------------------------------------------------------
+    # Tenant CRUD
+    # ------------------------------------------------------------------
+
+    def create(
+        self,
+        name: str,
+        config: Optional[Dict[str, Any]] = None,
+    ) -> tuple[TenantIdentity, str]:
+        """Create a new tenant. Returns (tenant, raw_api_key)."""
+        tenant_id = secrets.token_hex(12)
+        api_key = secrets.token_urlsafe(32)
+        api_key_hash = hashlib.sha256(api_key.encode()).hexdigest()
+
+        tenant = TenantIdentity(
+            tenant_id=tenant_id,
+            name=name,
+            config=config or {},
+            api_key_hash=api_key_hash,
+        )
+
+        self._conn.execute(
+            "INSERT INTO tenants (tenant_id, name, status, created_at, config, api_key_hash) "
+            "VALUES (?, ?, ?, ?, ?, ?)",
+            (
+                tenant.tenant_id,
+                tenant.name,
+                tenant.status,
+                tenant.created_at,
+                json.dumps(tenant.config),
+                tenant.api_key_hash,
+            ),
+        )
+        self._conn.commit()
+        return tenant, api_key
+
+    def get(self, tenant_id: str) -> Optional[TenantIdentity]:
+        row = self._conn.execute(
+            "SELECT * FROM tenants WHERE tenant_id = ? AND status != 'deleted'",
+            (tenant_id,),
+        ).fetchone()
+        return self._row_to_tenant(row) if row else None
+
+    def get_by_name(self, name: str) -> Optional[TenantIdentity]:
+        row = self._conn.execute(
+            "SELECT * FROM tenants WHERE name = ? AND status != 'deleted'",
+            (name,),
+        ).fetchone()
+        return self._row_to_tenant(row) if row else None
+
+    def authenticate(self, api_key: str) -> Optional[TenantIdentity]:
+        """Validate an API key and return the tenant if valid.
+
+        Scans all active tenants' key hashes. For large deployments, consider
+        a key-prefix lookup table instead.
+        """
+        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
+        row = self._conn.execute(
+            "SELECT * FROM tenants WHERE api_key_hash = ? AND status = 'active'",
+            (key_hash,),
+        ).fetchone()
+        return self._row_to_tenant(row) if row else None
+
+    def list_tenants(self, include_inactive: bool = False) -> List[TenantIdentity]:
+        if include_inactive:
+            rows = self._conn.execute(
+                "SELECT * FROM tenants ORDER BY created_at DESC"
+            ).fetchall()
+        else:
+            rows = self._conn.execute(
+                "SELECT * FROM tenants WHERE status = 'active' ORDER BY created_at DESC"
+            ).fetchall()
+        return [self._row_to_tenant(r) for r in rows]
+
+    def update_status(self, tenant_id: str, status: str) -> bool:
+        """Set tenant status: active, suspended, deleted."""
+        cur = self._conn.execute(
+            "UPDATE tenants SET status = ? WHERE tenant_id = ?",
+            (status, tenant_id),
+        )
+        self._conn.commit()
+        return cur.rowcount > 0
+
+    def rotate_key(self, tenant_id: str) -> Optional[str]:
+        """Generate a new API key for a tenant. Returns the new raw key."""
+        tenant = self.get(tenant_id)
+        if not tenant:
+            return None
+        new_key = secrets.token_urlsafe(32)
+        new_hash = hashlib.sha256(new_key.encode()).hexdigest()
+        self._conn.execute(
+            "UPDATE tenants SET api_key_hash = ? WHERE tenant_id = ?",
+            (new_hash, tenant_id),
+        )
+        self._conn.commit()
+        return new_key
+
+    def delete(self, tenant_id: str) -> bool:
+        """Soft-delete a tenant."""
+        return self.update_status(tenant_id, "deleted")
+
+    # ------------------------------------------------------------------
+    # Scoped API Keys
+    # ------------------------------------------------------------------
+
+    VALID_SCOPES = {"read", "read_write", "admin"}
+
+    def create_api_key(self, tenant_id: str, name: str, scope: str = "read_write",
+                       created_by: str = None, expires_hours: int = None) -> tuple:
+        """Create a scoped API key for a tenant.
+
+        Scopes:
+          - read: GET only (queries, searches, context reads)
+          - read_write: GET + POST/PUT (create nodes, ingest, run pipelines)
+          - admin: full access (delete, manage team, billing)
+
+        Returns (key_id, raw_key, scope).
+        """
+        if scope not in self.VALID_SCOPES:
+            raise ValueError(f"Invalid scope: {scope}. Must be one of: {self.VALID_SCOPES}")
+
+        key_id = secrets.token_hex(8)
+        raw_key = secrets.token_urlsafe(32)
+        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+        now = datetime.now(timezone.utc)
+
+        expires_at = None
+        if expires_hours:
+            from datetime import timedelta
+            expires_at = (now + timedelta(hours=expires_hours)).isoformat()
+
+        self._conn.execute(
+            "INSERT INTO api_keys (key_id, tenant_id, name, key_hash, scope, "
+            "created_by, created_at, expires_at, status) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'active')",
+            (key_id, tenant_id, name, key_hash, scope, created_by, now.isoformat(), expires_at),
+        )
+        self._conn.commit()
+        return key_id, raw_key, scope
+
+    def authenticate_api_key(self, raw_key: str):
+        """Validate a scoped API key. Returns (TenantIdentity, scope) or (None, None)."""
+        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+        row = self._conn.execute(
+            "SELECT ak.*, t.* FROM api_keys ak "
+            "JOIN tenants t ON ak.tenant_id = t.tenant_id "
+            "WHERE ak.key_hash = ? AND ak.status = 'active' AND t.status = 'active'",
+            (key_hash,),
+        ).fetchone()
+        if not row:
+            return None, None
+
+        # Check expiry
+        if row["expires_at"]:
+            now = datetime.now(timezone.utc).isoformat()
+            if now > row["expires_at"]:
+                return None, None
+
+        # Update last_used_at
+        self._conn.execute(
+            "UPDATE api_keys SET last_used_at = ? WHERE key_id = ?",
+            (datetime.now(timezone.utc).isoformat(), row["key_id"]),
+        )
+        self._conn.commit()
+
+        tenant = self._row_to_tenant(row)
+        return tenant, row["scope"]
+
+    def list_api_keys(self, tenant_id: str) -> List[Dict[str, Any]]:
+        """List all API keys for a tenant (without hashes)."""
+        rows = self._conn.execute(
+            "SELECT key_id, name, scope, created_by, created_at, last_used_at, "
+            "expires_at, status FROM api_keys WHERE tenant_id = ? ORDER BY created_at DESC",
+            (tenant_id,),
+        ).fetchall()
+        return [dict(r) for r in rows]
+
+    def revoke_api_key(self, key_id: str) -> bool:
+        """Revoke an API key."""
+        cur = self._conn.execute(
+            "UPDATE api_keys SET status = 'revoked' WHERE key_id = ?", (key_id,),
+        )
+        self._conn.commit()
+        return cur.rowcount > 0
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _row_to_tenant(self, row: sqlite3.Row) -> TenantIdentity:
+        return TenantIdentity(
+            tenant_id=row["tenant_id"],
+            name=row["name"],
+            status=row["status"],
+            created_at=row["created_at"],
+            config=json.loads(row["config"]) if row["config"] else {},
+            api_key_hash=row["api_key_hash"],
+        )
+
+    def close(self):
+        if self._conn:
+            self._conn.close()