mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
feat(capabilities): generate REST door with per-workspace rate limit
This commit is contained in:
parent
c696829568
commit
9c9ad97f8e
3 changed files with 135 additions and 0 deletions
1
surfsense_backend/app/capabilities/access/__init__.py
Normal file
1
surfsense_backend/app/capabilities/access/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Access doors (05): thin adapters that expose the 04 registry as REST/MCP/chat."""
|
||||
65
surfsense_backend/app/capabilities/access/rate_limit.py
Normal file
65
surfsense_backend/app/capabilities/access/rate_limit.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
"""Per-workspace rate limit for the capability doors (05).
|
||||
|
||||
A secondary abuse guard; the credit meter-gate (03c) is the primary control.
|
||||
Fixed-window over Redis (shared across workers) with a per-worker in-memory
|
||||
fallback when Redis is unavailable — mirroring the auth-endpoint limiter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from threading import Lock
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
||||
from app.config import config
|
||||
|
||||
CAPABILITY_RATE_LIMIT_PER_MINUTE = 120
|
||||
_WINDOW_SECONDS = 60
|
||||
_KEY_PREFIX = "surfsense:capability_rate_limit"
|
||||
|
||||
_redis = None
|
||||
_memory: dict[str, list[float]] = defaultdict(list)
|
||||
_memory_lock = Lock()
|
||||
|
||||
|
||||
def _redis_client():
|
||||
global _redis
|
||||
if _redis is None:
|
||||
import redis
|
||||
|
||||
_redis = redis.from_url(config.REDIS_APP_URL, decode_responses=True)
|
||||
return _redis
|
||||
|
||||
|
||||
def _incr_memory(key: str, window_seconds: int) -> int:
|
||||
now = time.monotonic()
|
||||
with _memory_lock:
|
||||
hits = [t for t in _memory[key] if now - t < window_seconds]
|
||||
hits.append(now)
|
||||
_memory[key] = hits
|
||||
return len(hits)
|
||||
|
||||
|
||||
def _incr(key: str, window_seconds: int) -> int:
|
||||
"""Increment the window counter for ``key`` and return the new count."""
|
||||
try:
|
||||
client = _redis_client()
|
||||
count = int(client.incr(key))
|
||||
if count == 1:
|
||||
client.expire(key, window_seconds)
|
||||
return count
|
||||
except Exception:
|
||||
return _incr_memory(key, window_seconds)
|
||||
|
||||
|
||||
async def enforce_capability_rate_limit(request: Request) -> None:
|
||||
"""Cap requests per workspace per minute; raise 429 when exceeded."""
|
||||
workspace_id = request.path_params.get("workspace_id")
|
||||
count = _incr(f"{_KEY_PREFIX}:{workspace_id}", _WINDOW_SECONDS)
|
||||
if count > CAPABILITY_RATE_LIMIT_PER_MINUTE:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail="Rate limit exceeded for this workspace. Try again shortly.",
|
||||
)
|
||||
69
surfsense_backend/app/capabilities/access/rest.py
Normal file
69
surfsense_backend/app/capabilities/access/rest.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
"""Generate the REST door from the capability registry (05).
|
||||
|
||||
One typed ``POST`` per verb; each runs the same thin adapter:
|
||||
authn -> workspace authz -> meter-gate -> executor -> charge -> typed output.
|
||||
"""
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.auth.context import AuthContext
|
||||
from app.capabilities.access.rate_limit import enforce_capability_rate_limit
|
||||
from app.capabilities.billing import charge_capability, gate_capability
|
||||
from app.capabilities.store import all_capabilities
|
||||
from app.capabilities.types import Capability, CapabilityContext
|
||||
from app.db import get_async_session
|
||||
from app.services.web_crawl_credit_service import InsufficientCreditsError
|
||||
from app.users import get_auth_context
|
||||
from app.utils.rbac import check_workspace_access
|
||||
|
||||
|
||||
def build_capabilities_router(
|
||||
capabilities: list[Capability] | None = None,
|
||||
) -> APIRouter:
|
||||
"""Emit one typed route per verb (defaults to the whole registry)."""
|
||||
router = APIRouter(tags=["capabilities"])
|
||||
caps = capabilities if capabilities is not None else all_capabilities()
|
||||
for capability in caps:
|
||||
_register_verb(router, capability)
|
||||
return router
|
||||
|
||||
|
||||
def _register_verb(router: APIRouter, capability: Capability) -> None:
|
||||
input_model = capability.input_schema
|
||||
output_model = capability.output_schema
|
||||
unit = capability.billing_unit
|
||||
executor = capability.executor
|
||||
|
||||
async def endpoint(
|
||||
workspace_id: int,
|
||||
payload: input_model,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
auth: AuthContext = Depends(get_auth_context),
|
||||
):
|
||||
await check_workspace_access(session, auth, workspace_id)
|
||||
ctx = CapabilityContext(session=session, workspace_id=workspace_id)
|
||||
try:
|
||||
await gate_capability(payload, unit, ctx)
|
||||
except InsufficientCreditsError as exc:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_402_PAYMENT_REQUIRED,
|
||||
detail={
|
||||
"error_code": "insufficient_credits",
|
||||
"message": str(exc),
|
||||
"balance_micros": exc.balance_micros,
|
||||
"required_micros": exc.required_micros,
|
||||
},
|
||||
) from exc
|
||||
output = await executor(payload)
|
||||
await charge_capability(output, unit, ctx)
|
||||
return output
|
||||
|
||||
router.add_api_route(
|
||||
f"/workspaces/{{workspace_id}}/capabilities/{capability.name}",
|
||||
endpoint,
|
||||
methods=["POST"],
|
||||
response_model=output_model,
|
||||
name=f"capability:{capability.name}",
|
||||
dependencies=[Depends(enforce_capability_rate_limit)],
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue