feat(ai): durcissement phase F (#79) - rate limit, redaction, OpenAPI/MCP, E2E

This commit is contained in:
2026-09-11 21:56:50 -04:00
parent 88eecd7671
commit 9dce341bc8
18 changed files with 854 additions and 27 deletions
+16
View File
@@ -17,11 +17,22 @@ from backend.tools.context import (
ToolMode,
ToolNotFoundError,
ToolPermissionError,
ToolRateLimitError,
ToolRisk,
ToolScope,
ToolValidationError,
resolve_safe_path,
)
from backend.tools.ratelimit import (
check_and_record as check_tool_rate_limit,
)
from backend.tools.ratelimit import (
get_status as get_tool_rate_limit_status,
)
from backend.tools.ratelimit import (
reset as reset_tool_rate_limit,
)
from backend.tools.redaction import redact_payload
from backend.tools.registry import (
ToolSpec,
call_tool,
@@ -39,15 +50,20 @@ __all__ = [
"ToolMode",
"ToolNotFoundError",
"ToolPermissionError",
"ToolRateLimitError",
"ToolResult",
"ToolRisk",
"ToolScope",
"ToolSpec",
"ToolValidationError",
"call_tool",
"check_tool_rate_limit",
"get_tool",
"get_tool_rate_limit_status",
"get_tool_schemas",
"list_tools",
"redact_payload",
"reset_tool_rate_limit",
"resolve_safe_path",
"tool",
]
+30
View File
@@ -102,6 +102,21 @@ class ToolConfirmationRequired(ToolError):
return payload
class ToolRateLimitError(ToolError):
"""Raised when an identity exceeds its tool-call rate limit (Phase F)."""
code = "rate_limited"
def __init__(self, tool: str, retry_after: int = 0, message: str | None = None):
super().__init__(
message or f"Rate limit exceeded for tool '{tool}'",
code="rate_limited",
details={"tool": tool, "retry_after": retry_after},
)
self.tool = tool
self.retry_after = retry_after
@dataclass
class ToolContext:
"""Carries the caller identity and execution options for a tool call.
@@ -126,6 +141,21 @@ class ToolContext:
def username(self) -> str:
return self.user.get("username", "unknown")
@property
def rate_limit_identity(self) -> str:
"""Identity used by the per-token/per-tool rate limiter.
Prefers the JWT id (``_token_jti``, attached by the auth middleware),
then an explicit ``token_id`` in :attr:`metadata`, then the user id and
finally the username. This yields per-token limiting when a token id is
available and per-account limiting otherwise.
"""
token_id = self.metadata.get("token_id") or self.user.get("_token_jti")
if token_id:
return f"token:{token_id}"
user_id = self.user.get("id") or self.user.get("username")
return f"user:{user_id or 'anonymous'}"
def has_vault_access(self, vault: str) -> bool:
"""Return True if the caller may access *vault*."""
return bool(vault) and check_vault_access(vault, self.user)
+153
View File
@@ -0,0 +1,153 @@
"""Rate limiting for the AI tool layer (Phase F).
Complements the IP-based login limiter (``backend.ratelimit``) with a
per-identity, per-tool sliding-window limiter applied to every tool call,
whether it originates from the in-app assistant or the MCP server.
Two counters are maintained per identity:
* a **global** counter (all tools combined), capped by
``OBSIGATE_TOOL_RATE_LIMIT``;
* a **per-tool** counter, capped by ``OBSIGATE_TOOL_RATE_LIMIT_PER_TOOL``
(defaults to the global cap) so a single expensive tool cannot consume the
whole budget.
The window length is ``OBSIGATE_TOOL_RATE_WINDOW`` seconds (default 60).
The identity is the token JTI when available (``_token_jti`` attached by the
auth middleware), otherwise the user id/username. This gives "per token / per
tool" limiting while remaining meaningful for anonymous/disabled-auth mode.
"""
from __future__ import annotations
import logging
import os
import time
from collections import deque
from typing import Any
logger = logging.getLogger("obsigate.tools.ratelimit")
# --- Configuration (read at call time so tests can monkeypatch env) ---
DEFAULT_RATE_LIMIT = 60
DEFAULT_RATE_WINDOW = 60
def _env_int(name: str, default: int) -> int:
try:
return int(os.environ.get(name, str(default)))
except (TypeError, ValueError):
return default
def _global_limit() -> int:
return _env_int("OBSIGATE_TOOL_RATE_LIMIT", DEFAULT_RATE_LIMIT)
def _per_tool_limit() -> int:
return _env_int("OBSIGATE_TOOL_RATE_LIMIT_PER_TOOL", _global_limit())
def _window() -> int:
return max(1, _env_int("OBSIGATE_TOOL_RATE_WINDOW", DEFAULT_RATE_WINDOW))
# --- In-memory store: {counter_key: deque[timestamp]} ---
_calls: dict[str, deque] = {}
def _identity_key(identity: str | None) -> str:
return identity or "anonymous"
def _counter_key(identity: str | None, tool: str | None) -> str:
base = _identity_key(identity)
return f"{base}\x00{tool}" if tool else base
def _prune(entries: deque, now: float, window: int) -> None:
cutoff = now - window
while entries and entries[0] <= cutoff:
entries.popleft()
def check_and_record(identity: str | None, tool: str, *, now: float | None = None) -> tuple[bool, int]:
"""Record one call and report whether it is allowed.
Args:
identity: Rate-limit identity (token JTI or username).
tool: Tool name (used for the per-tool counter).
now: Optional timestamp override (tests).
Returns:
``(allowed, retry_after)``. When ``allowed`` is False, ``retry_after``
is the number of seconds until the oldest call leaves the window.
"""
now = time.time() if now is None else now
window = _window()
global_limit = _global_limit()
tool_limit = _per_tool_limit()
global_entries = _calls.setdefault(_counter_key(identity, None), deque())
tool_entries = _calls.setdefault(_counter_key(identity, tool), deque())
_prune(global_entries, now, window)
_prune(tool_entries, now, window)
if len(global_entries) >= global_limit or len(tool_entries) >= tool_limit:
oldest = min(
global_entries[0] if len(global_entries) >= global_limit else now,
tool_entries[0] if len(tool_entries) >= tool_limit else now,
)
retry_after = max(1, int(oldest + window - now) + 1)
logger.warning(f"Tool rate limit exceeded for '{_identity_key(identity)}' on '{tool}'")
return False, retry_after
global_entries.append(now)
tool_entries.append(now)
return True, 0
def remaining(identity: str | None, tool: str | None = None) -> int:
"""Return the number of calls still allowed in the current window."""
now = time.time()
window = _window()
if tool:
entries = _calls.get(_counter_key(identity, tool))
if entries is None:
return _per_tool_limit()
_prune(entries, now, window)
return max(0, _per_tool_limit() - len(entries))
entries = _calls.get(_counter_key(identity, None))
if entries is None:
return _global_limit()
_prune(entries, now, window)
return max(0, _global_limit() - len(entries))
def reset(identity: str | None = None) -> None:
"""Clear rate-limit state (all identities, or a single one)."""
if identity is None:
_calls.clear()
return
prefix = f"{_identity_key(identity)}\x00"
for key in [k for k in _calls if k == _identity_key(identity) or k.startswith(prefix)]:
del _calls[key]
def get_status(identity: str | None = None) -> dict[str, Any]:
"""Diagnostic snapshot of the limiter (for tests / admin tooling)."""
if identity is None:
return {
"tracked_identities": len({k.split("\x00", 1)[0] for k in _calls}),
"global_limit": _global_limit(),
"per_tool_limit": _per_tool_limit(),
"window_seconds": _window(),
}
return {
"identity": _identity_key(identity),
"remaining_global": remaining(identity),
"global_limit": _global_limit(),
"per_tool_limit": _per_tool_limit(),
"window_seconds": _window(),
}
+39
View File
@@ -0,0 +1,39 @@
"""Secret redaction for tool results (Phase F).
Tool results are fed back to the LLM (in-app agent loop) or returned to an
external MCP client, so they must never leak credentials. ``read_file`` and
``read_file_raw`` already redact through the shared file service, but other
tools (``diff_backup``, ``search_*``) return content that has not been through
the redactor. This module applies :func:`backend.secret_redactor.redact` to
every string in a tool payload, recursively.
"""
from __future__ import annotations
from typing import Any
from backend.secret_redactor import redact
_MAX_DEPTH = 12
def redact_payload(payload: Any, _depth: int = 0) -> Any:
"""Return a copy of *payload* with every string secret-redacted.
Walks dicts, lists and tuples; scalars are returned unchanged. Strings are
run through :func:`backend.secret_redactor.redact`. A recursion cap avoids
pathological/cyclic structures (tool results are JSON-serialisable, so
cycles should not occur, but the guard keeps this safe).
"""
if _depth > _MAX_DEPTH:
return payload
if isinstance(payload, str):
redacted, count = redact(payload)
return redacted if count else payload
if isinstance(payload, dict):
return {key: redact_payload(value, _depth + 1) for key, value in payload.items()}
if isinstance(payload, list):
return [redact_payload(item, _depth + 1) for item in payload]
if isinstance(payload, tuple):
return tuple(redact_payload(item, _depth + 1) for item in payload)
return payload
+12 -1
View File
@@ -26,10 +26,13 @@ from backend.tools.context import (
ToolError,
ToolNotFoundError,
ToolPermissionError,
ToolRateLimitError,
ToolRisk,
ToolScope,
ToolValidationError,
)
from backend.tools.ratelimit import check_and_record
from backend.tools.redaction import redact_payload
from backend.tools.schemas import ToolResult
logger = logging.getLogger("obsigate.tools.registry")
@@ -186,6 +189,12 @@ def call_tool(
_audit(ctx, spec, arguments, ok=False, error=error.code)
raise error from e
allowed, retry_after = check_and_record(ctx.rate_limit_identity, spec.name)
if not allowed:
rate_error = ToolRateLimitError(spec.name, retry_after)
_audit(ctx, spec, arguments, ok=False, error=rate_error.code)
raise rate_error
vault = getattr(params, "vault", None)
if spec.requires_vault and vault:
try:
@@ -223,4 +232,6 @@ def call_tool(
duration_ms = round((time.perf_counter() - started) * 1000, 2)
logger.debug(f"Tool '{name}' ok in {duration_ms}ms")
_audit(ctx, spec, arguments, ok=True)
return ToolResult(ok=True, data=data)
# Never forward raw secrets to the model (defense in depth: some tools —
# diffs, search snippets — return content that was not pre-redacted).
return ToolResult(ok=True, data=redact_payload(data))
+4 -2
View File
@@ -11,6 +11,7 @@ confirmation mechanism (two-step propose/apply).
from __future__ import annotations
import logging
import os
from typing import Any
from backend.indexer import get_backlinks as _get_backlinks
@@ -90,8 +91,9 @@ from backend.tools.schemas import (
logger = logging.getLogger("obsigate.tools.service")
# Maximum file size returned by ``read_file`` (bytes).
TOOL_MAX_READ_BYTES = 200_000
# Maximum file size returned by ``read_file`` (bytes). Quota configurable via
# ``BOOKSLM_MAX_TOOL_READ_BYTES``.
TOOL_MAX_READ_BYTES = int(os.environ.get("BOOKSLM_MAX_TOOL_READ_BYTES", "200000"))
# ── C1. Vaults / navigation ────────────────────────────────────────────────