refactor: harden rate limit diagnostics and cooldowns

This commit is contained in:
rookiestar28
2026-03-19 22:56:02 +08:00
parent 7519eb0efd
commit 1b104b170b
20 changed files with 899 additions and 148 deletions
+16 -4
View File
@@ -12,7 +12,7 @@ try:
from ..services.automation_composer import AutomationComposerService
from ..services.planner import PlannerService
from ..services.planner_registry import get_planner_registry
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.reasoning_redaction import (
audit_reasoning_reveal,
resolve_reasoning_reveal,
@@ -26,7 +26,7 @@ except ImportError:
from services.automation_composer import AutomationComposerService
from services.planner import PlannerService
from services.planner_registry import get_planner_registry
from services.rate_limit import check_rate_limit
from services.rate_limit import build_rate_limit_response, check_rate_limit
from services.reasoning_redaction import (
audit_reasoning_reveal,
resolve_reasoning_reveal,
@@ -90,7 +90,13 @@ class AssistHandlers:
if not authorized:
return web.json_response({"error": "Unauthorized"}, status=401)
if not check_rate_limit(request, "admin"):
return web.json_response({"error": "Rate limit exceeded"}, status=429)
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=False,
)
return None
async def _parse_json_body(
@@ -557,7 +563,13 @@ class AssistHandlers:
return web.json_response({"error": "Unauthorized"}, status=401)
if not check_rate_limit(request, "admin"):
return web.json_response({"error": "Rate limit exceeded"}, status=429)
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=False,
)
try:
data = await request.json()
+16 -4
View File
@@ -25,7 +25,7 @@ try:
from ..services.bridge_handshake import verify_handshake
from ..services.execution_budgets import BudgetExceededError
from ..services.idempotency_store import IdempotencyStore
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.sidecar.auth import is_bridge_enabled, require_bridge_auth
from ..services.sidecar.bridge_contract import (
BRIDGE_ENDPOINTS,
@@ -43,7 +43,7 @@ except ImportError:
from services.bridge_handshake import verify_handshake
from services.execution_budgets import BudgetExceededError
from services.idempotency_store import IdempotencyStore
from services.rate_limit import check_rate_limit
from services.rate_limit import build_rate_limit_response, check_rate_limit
from services.sidecar.auth import is_bridge_enabled, require_bridge_auth
from services.sidecar.bridge_contract import (
BRIDGE_ENDPOINTS,
@@ -255,7 +255,13 @@ class BridgeHandlers:
scope=BridgeScope.JOB_SUBMIT.value,
details={"reason": "rate_limit"},
)
return web.json_response({"error": "Rate limit exceeded"}, status=429)
return build_rate_limit_response(
request,
"bridge",
web_module=web,
error="Rate limit exceeded",
include_ok=False,
)
# Parse payload
try:
@@ -484,7 +490,13 @@ class BridgeHandlers:
scope=BridgeScope.DELIVERY.value,
details={"reason": "rate_limit"},
)
return web.json_response({"error": "Rate limit exceeded"}, status=429)
return build_rate_limit_response(
request,
"bridge",
web_module=web,
error="Rate limit exceeded",
include_ok=False,
)
# Parse payload
try:
+23 -4
View File
@@ -24,7 +24,7 @@ if __package__ and "." in __package__:
get_checkpoint,
list_checkpoints,
)
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_payload, check_rate_limit
from ..services.request_ip import get_client_ip
else: # pragma: no cover (test-only import mode)
from models.schemas import MAX_BODY_SIZE # type: ignore
@@ -35,7 +35,10 @@ else: # pragma: no cover (test-only import mode)
get_checkpoint,
list_checkpoints,
)
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_payload,
check_rate_limit,
)
from services.request_ip import get_client_ip # type: ignore
@@ -105,7 +108,15 @@ async def list_checkpoints_handler(request: web.Request) -> web.Response:
raise RuntimeError("aiohttp not available")
if not check_rate_limit(request, "admin"):
return _json_resp({"ok": False, "error": "rate_limit_exceeded"}, 429)
return _json_resp(
build_rate_limit_payload(
request,
"admin",
error="rate_limit_exceeded",
include_ok=True,
),
429,
)
allowed, error = require_admin_token(request)
if not allowed:
@@ -136,7 +147,15 @@ async def create_checkpoint_handler(request: web.Request) -> web.Response:
raise RuntimeError("aiohttp not available")
if not check_rate_limit(request, "admin"):
return _json_resp({"ok": False, "error": "rate_limit_exceeded"}, 429)
return _json_resp(
build_rate_limit_payload(
request,
"admin",
error="rate_limit_exceeded",
include_ok=True,
),
429,
)
# Body Size Check
if request.content_length and request.content_length > MAX_BODY_SIZE:
+32 -12
View File
@@ -107,11 +107,11 @@ except Exception:
"services.llm_client",
("LLMClient",),
)
(check_rate_limit,) = import_attrs_dual(
(check_rate_limit, build_rate_limit_response) = import_attrs_dual(
__package__,
"..services.rate_limit",
"services.rate_limit",
("check_rate_limit",),
("check_rate_limit", "build_rate_limit_response"),
)
(get_client_ip,) = import_attrs_dual(
__package__,
@@ -281,8 +281,12 @@ async def config_get_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
token_info = resolve_token_info(request)
@@ -359,8 +363,12 @@ async def llm_models_handler(request: web.Request) -> web.Response:
raise RuntimeError("aiohttp not available")
# S17: Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
token_info = resolve_token_info(request)
@@ -578,8 +586,12 @@ async def config_put_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
# R99/S46: resolve identity context for non-repudiation audits.
@@ -776,8 +788,12 @@ async def llm_test_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
token_info = resolve_token_info(request)
@@ -997,8 +1013,12 @@ async def llm_chat_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
# NOTE: Keep this server-side. Connector cannot access UI-stored secrets directly.
+11 -4
View File
@@ -28,7 +28,7 @@ if __package__ and "." in __package__:
from ..services.connector_installation_registry import (
get_connector_installation_registry,
)
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.tenant_context import TenantBoundaryError, request_tenant_scope
else: # pragma: no cover
from services.access_control import require_admin_token # type: ignore
@@ -36,7 +36,10 @@ else: # pragma: no cover
from services.connector_installation_registry import ( # type: ignore
get_connector_installation_registry,
)
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.tenant_context import ( # type: ignore
TenantBoundaryError,
request_tenant_scope,
@@ -62,8 +65,12 @@ logger = logging.getLogger("ComfyUI-OpenClaw.api.connector_contracts")
def _require_admin(request) -> Optional[web.Response]:
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
allowed, err = require_admin_token(request)
if not allowed:
+17 -10
View File
@@ -30,7 +30,7 @@ if __package__ and "." in __package__:
from ..services.job_events import get_job_event_store
from ..services.management_query import normalize_cursor_limit
from ..services.metrics import metrics
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.reasoning_redaction import (
audit_reasoning_reveal,
resolve_reasoning_reveal,
@@ -43,7 +43,10 @@ else: # pragma: no cover
from services.job_events import get_job_event_store # type: ignore
from services.management_query import normalize_cursor_limit # type: ignore
from services.metrics import metrics # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.reasoning_redaction import ( # type: ignore
audit_reasoning_reveal,
resolve_reasoning_reveal,
@@ -112,10 +115,12 @@ async def events_stream_handler(request: web.Request) -> web.StreamResponse:
# Rate limit
if not check_rate_limit(request, "events"):
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"events",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
# Access control (same as logs/tail)
@@ -223,10 +228,12 @@ async def events_poll_handler(request: web.Request) -> web.Response:
# Rate limit
if not check_rate_limit(request, "events"):
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"events",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
# Access control
+17 -6
View File
@@ -24,7 +24,7 @@ if __package__ and "." in __package__:
get_model_inventory_snapshot,
run_preflight_check,
)
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.request_ip import get_client_ip
else: # pragma: no cover (test-only import mode)
from models.schemas import MAX_BODY_SIZE # type: ignore
@@ -34,7 +34,10 @@ else: # pragma: no cover (test-only import mode)
get_model_inventory_snapshot,
run_preflight_check,
)
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.request_ip import get_client_ip # type: ignore
# R98: Endpoint Metadata
@@ -102,8 +105,12 @@ async def preflight_handler(request: web.Request) -> web.Response:
# Rate limit: admin-grade endpoint (inventory leak + CPU cost)
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
# Body Size Check
@@ -176,8 +183,12 @@ async def inventory_handler(request: web.Request) -> web.Response:
# Rate Limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
# Admin boundary (localhost convenience mode if no token configured)
+8 -6
View File
@@ -281,11 +281,11 @@ if web is not None:
"update_experiment_handler",
),
)
(check_rate_limit,) = import_attrs_dual(
(check_rate_limit, build_rate_limit_response) = import_attrs_dual(
__package__,
"..services.rate_limit",
"services.rate_limit",
("check_rate_limit",),
("check_rate_limit", "build_rate_limit_response"),
)
(redact_text,) = import_attrs_dual(
__package__,
@@ -502,10 +502,12 @@ async def logs_tail_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "logs"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"logs",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
try:
+11 -6
View File
@@ -30,7 +30,7 @@ if __package__ and "." in __package__:
from ..services.audit import emit_audit_event
from ..services.csrf_protection import require_same_origin_if_no_token
from ..services.metrics import metrics
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.request_ip import get_client_ip
from ..services.runtime_config import get_admin_token, is_loopback_client
from ..services.secret_store import get_secret_store
@@ -41,7 +41,10 @@ else: # pragma: no cover (test-only import mode)
from services.audit import emit_audit_event # type: ignore
from services.csrf_protection import require_same_origin_if_no_token # type: ignore
from services.metrics import metrics # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.request_ip import get_client_ip # type: ignore
from services.runtime_config import get_admin_token # type: ignore
from services.runtime_config import is_loopback_client
@@ -120,10 +123,12 @@ def _rate_limit_admin(request: web.Request) -> Optional[web.Response]:
if check_rate_limit(request, "admin"):
return None
metrics.increment("rate_limit_exceeded")
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
+11 -4
View File
@@ -41,11 +41,14 @@ except ImportError: # pragma: no cover
if __package__ and "." in __package__:
from ..services.access_control import require_admin_token
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.security_doctor import run_security_doctor
else: # pragma: no cover (test-only)
from services.access_control import require_admin_token # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.security_doctor import run_security_doctor # type: ignore
# R98: Endpoint Metadata
@@ -90,8 +93,12 @@ async def security_doctor_handler(request: web.Request) -> web.Response:
# S17: Rate limit
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
# Admin boundary
+11 -6
View File
@@ -22,7 +22,7 @@ if __package__ and "." in __package__:
require_observability_access,
resolve_token_info,
)
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.templates import get_template_service
from ..services.tenant_context import TenantBoundaryError, request_tenant_scope
else: # pragma: no cover (test-only import mode)
@@ -30,7 +30,10 @@ else: # pragma: no cover (test-only import mode)
require_observability_access,
resolve_token_info,
)
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.templates import get_template_service # type: ignore
from services.tenant_context import ( # type: ignore
TenantBoundaryError,
@@ -108,10 +111,12 @@ async def templates_list_handler(request: web.Request) -> web.Response:
# Reuse the admin bucket to avoid unbounded enumeration from remote callers.
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "Rate limit exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
try:
+8 -7
View File
@@ -23,13 +23,13 @@ except ImportError:
try:
from ..models.schemas import MAX_BODY_SIZE, WebhookJobRequest
from ..services.metrics import metrics
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.trace import get_effective_trace_id
from ..services.webhook_auth import get_auth_summary, require_auth
except ImportError:
from models.schemas import MAX_BODY_SIZE, WebhookJobRequest
from services.metrics import metrics
from services.rate_limit import check_rate_limit
from services.rate_limit import build_rate_limit_response, check_rate_limit
from services.trace import get_effective_trace_id
from services.webhook_auth import get_auth_summary, require_auth
@@ -71,11 +71,12 @@ async def webhook_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "webhook"):
metrics.inc("webhook_denied")
return create_error_response(
message="Rate limit exceeded",
code=ErrorCode.RATE_LIMIT_EXCEEDED,
status=429,
detail={"retry_after": "60"},
return build_rate_limit_response(
request,
"webhook",
web_module=web,
error="Rate limit exceeded",
include_ok=True,
)
try:
+11 -6
View File
@@ -23,7 +23,7 @@ if __package__ and "." in __package__:
from ..services.job_events import JobEventType, get_job_event_store
from ..services.metrics import metrics
from ..services.queue_submit import submit_prompt
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.templates import get_template_service
from ..services.trace import get_effective_trace_id
from ..services.trace_store import trace_store
@@ -41,7 +41,10 @@ else: # pragma: no cover (test-only import mode)
from services.job_events import JobEventType, get_job_event_store # type: ignore
from services.metrics import metrics # type: ignore
from services.queue_submit import submit_prompt # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.templates import get_template_service # type: ignore
from services.trace import get_effective_trace_id # type: ignore
from services.trace_store import trace_store # type: ignore
@@ -110,10 +113,12 @@ async def webhook_submit_handler(request: web.Request) -> web.Response:
# S17: Rate Limit
if not check_rate_limit(request, "webhook"):
metrics.inc("webhook_denied")
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"webhook",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
try:
+11 -6
View File
@@ -39,7 +39,7 @@ if __package__ and "." in __package__:
from ..models.schemas import MAX_BODY_SIZE, WebhookJobRequest
from ..services.execution_budgets import BudgetExceededError, check_render_size
from ..services.metrics import metrics
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
from ..services.templates import get_template_service
from ..services.trace import get_effective_trace_id
from ..services.webhook_auth import require_auth
@@ -51,7 +51,10 @@ else: # pragma: no cover (test-only import mode)
check_render_size,
)
from services.metrics import metrics # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
from services.templates import get_template_service # type: ignore
from services.trace import get_effective_trace_id # type: ignore
from services.webhook_auth import require_auth # type: ignore
@@ -106,10 +109,12 @@ async def webhook_validate_handler(request: web.Request) -> web.Response:
# S17: Rate limit (same bucket as submit)
if not check_rate_limit(request, "webhook"):
metrics.inc("webhook_denied")
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"},
status=429,
headers={"Retry-After": "60"},
return build_rate_limit_response(
request,
"webhook",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
# S2: Content-Type + body size
+82 -9
View File
@@ -43,12 +43,23 @@ class CooldownEntry:
model: Optional[str]
reason: str
until: float # Unix timestamp when cooldown expires
reason_code: str = "provider_unknown"
bucket: str = "provider_unknown"
retry_after_sec: Optional[int] = None
def is_active(self) -> bool:
"""Check if cooldown is still active."""
return time.time() < self.until
@dataclass(frozen=True)
class CooldownDecision:
category: "ErrorCategory"
retry_after_sec: Optional[int]
reason_code: str
bucket: str
class FailoverState:
"""
Manages cooldown state persistence.
@@ -158,7 +169,15 @@ class FailoverState:
return provider
def set_cooldown(
self, provider: str, model: Optional[str], reason: str, duration_sec: float
self,
provider: str,
model: Optional[str],
reason: str,
duration_sec: float,
*,
reason_code: Optional[str] = None,
bucket: Optional[str] = None,
retry_after_sec: Optional[int] = None,
) -> None:
"""
Set a cooldown for a provider/model.
@@ -173,7 +192,13 @@ class FailoverState:
until = time.time() + duration_sec
self.cooldowns[key] = CooldownEntry(
provider=provider, model=model, reason=reason, until=until
provider=provider,
model=model,
reason=reason,
until=until,
reason_code=reason_code or reason,
bucket=bucket or "provider_unknown",
retry_after_sec=retry_after_sec,
)
self._save()
logger.info(f"Set cooldown for {key}: {reason} (until {until})")
@@ -317,9 +342,9 @@ def reset_failover_state(*, flush: bool = False) -> None:
_failover_state = None
def classify_error(
def classify_cooldown(
error: Exception, status_code: Optional[int] = None
) -> Tuple[ErrorCategory, Optional[int]]:
) -> CooldownDecision:
"""
Classify an error into a failover category and extract retry-after.
@@ -328,7 +353,7 @@ def classify_error(
status_code: Optional HTTP status code (may be in exception).
Returns:
Tuple of (ErrorCategory, retry_after_seconds or None)
Structured cooldown classification.
"""
# R14/R37: Check if error is ProviderHTTPError
try:
@@ -338,42 +363,90 @@ def classify_error(
status_code = error.status_code
retry_after = error.retry_after
else:
retry_after = None
retry_after = getattr(error, "retry_after", None)
except ImportError:
retry_after = None
retry_after = getattr(error, "retry_after", None)
error_str = str(error).lower()
reason_code = "provider_unknown"
bucket = "provider_unknown"
# Status code-based classification
if status_code:
if status_code == 401 or status_code == 403:
category = ErrorCategory.AUTH
reason_code = "provider_auth_failed"
bucket = "provider_auth"
elif status_code == 402:
category = ErrorCategory.BILLING
reason_code = "provider_billing_required"
bucket = "provider_quota"
elif status_code == 429:
# Distinguish rate limit vs billing
if "quota" in error_str or "billing" in error_str:
if (
"quota" in error_str
or "billing" in error_str
or "insufficient_quota" in error_str
or "insufficient quota" in error_str
):
category = ErrorCategory.BILLING
reason_code = "provider_quota_exceeded"
bucket = "provider_quota"
else:
category = ErrorCategory.RATE_LIMIT
if retry_after is not None:
reason_code = "provider_retry_after"
else:
reason_code = "provider_rate_limited"
bucket = "provider_cooldown"
elif status_code == 400 or status_code == 422:
category = ErrorCategory.INVALID_REQUEST
reason_code = "provider_invalid_request"
bucket = "provider_invalid_request"
else:
category = ErrorCategory.UNKNOWN
reason_code = f"provider_http_{status_code}"
bucket = "provider_unknown"
else:
# Exception type-based classification
if "timeout" in error_str or "timed out" in error_str:
category = ErrorCategory.TIMEOUT
reason_code = "provider_timeout"
bucket = "provider_cooldown"
elif "unauthorized" in error_str or "forbidden" in error_str:
category = ErrorCategory.AUTH
reason_code = "provider_auth_failed"
bucket = "provider_auth"
elif "rate limit" in error_str or "too many requests" in error_str:
category = ErrorCategory.RATE_LIMIT
reason_code = (
"provider_retry_after"
if retry_after is not None
else "provider_rate_limited"
)
bucket = "provider_cooldown"
elif "quota" in error_str or "insufficient" in error_str:
category = ErrorCategory.BILLING
reason_code = "provider_quota_exceeded"
bucket = "provider_quota"
else:
category = ErrorCategory.UNKNOWN
reason_code = "provider_unknown"
bucket = "provider_unknown"
return category, retry_after
return CooldownDecision(
category=category,
retry_after_sec=retry_after,
reason_code=reason_code,
bucket=bucket,
)
def classify_error(
error: Exception, status_code: Optional[int] = None
) -> Tuple[ErrorCategory, Optional[int]]:
decision = classify_cooldown(error, status_code)
return decision.category, decision.retry_after_sec
def should_retry(category: ErrorCategory) -> bool:
+18 -2
View File
@@ -278,6 +278,7 @@ class LLMClient:
try:
from ..services.failover import (
ErrorCategory,
classify_cooldown,
classify_error,
get_cooldown_duration,
get_failover_state,
@@ -287,6 +288,7 @@ class LLMClient:
except ImportError:
from services.failover import (
ErrorCategory,
classify_cooldown,
classify_error,
get_cooldown_duration,
get_failover_state,
@@ -314,6 +316,7 @@ class LLMClient:
return {
"ErrorCategory": ErrorCategory,
"classify_cooldown": classify_cooldown,
"classify_error": classify_error,
"get_cooldown_duration": get_cooldown_duration,
"should_failover": should_failover,
@@ -454,6 +457,7 @@ class LLMClient:
) -> Dict[str, Any]:
"""R130 phase 2: execute failover/retry loop against prepared candidates."""
ErrorCategory = phase["ErrorCategory"]
classify_cooldown = phase["classify_cooldown"]
classify_error = phase["classify_error"]
get_cooldown_duration = phase["get_cooldown_duration"]
should_failover = phase["should_failover"]
@@ -620,6 +624,7 @@ class LLMClient:
except Exception as e:
candidate_last_error = e
status_code = self._extract_status_code(e)
cooldown_decision = classify_cooldown(e, status_code)
error_category, retry_after = classify_error(e, status_code)
logger.error(
f"Request failed for {provider}/{self.model}: {e} "
@@ -634,6 +639,8 @@ class LLMClient:
"model": self.model,
"candidate_index": candidate_idx,
"category": error_category.value,
"cooldown_bucket": cooldown_decision.bucket,
"reason_code": cooldown_decision.reason_code,
"status_code": status_code,
"error_type": type(e).__name__,
"trace_id": trace_id,
@@ -666,11 +673,17 @@ class LLMClient:
error_category, retry_after_override=retry_after
)
failover_state.set_cooldown(
provider, model, error_category.value, duration
provider,
model,
cooldown_decision.reason_code,
duration,
reason_code=cooldown_decision.reason_code,
bucket=cooldown_decision.bucket,
retry_after_sec=retry_after,
)
logger.warning(
f"Failover triggered for {provider}/{model}: "
f"{error_category.value} (cooldown: {duration}s)"
f"{cooldown_decision.reason_code} (cooldown: {duration}s)"
)
emit_structured_log(
logger,
@@ -680,7 +693,10 @@ class LLMClient:
"provider": provider,
"model": model,
"category": error_category.value,
"cooldown_bucket": cooldown_decision.bucket,
"reason_code": cooldown_decision.reason_code,
"cooldown_sec": duration,
"retry_after_sec": retry_after,
"trace_id": trace_id,
},
)
+11 -4
View File
@@ -21,10 +21,13 @@ except ImportError:
if __package__ and "." in __package__:
from ..services.access_control import require_admin_token
from ..services.rate_limit import check_rate_limit
from ..services.rate_limit import build_rate_limit_response, check_rate_limit
else: # pragma: no cover (test-only import mode)
from services.access_control import require_admin_token # type: ignore
from services.rate_limit import check_rate_limit # type: ignore
from services.rate_limit import ( # type: ignore
build_rate_limit_response,
check_rate_limit,
)
# R98: Endpoint Metadata
if __package__ and "." in __package__:
@@ -374,8 +377,12 @@ def _require_admin(request: web.Request) -> Optional[web.Response]:
auth + rate limit gates to avoid remote abuse and queue-flood vectors.
"""
if not check_rate_limit(request, "admin"):
return web.json_response(
{"ok": False, "error": "rate_limit_exceeded"}, status=429
return build_rate_limit_response(
request,
"admin",
web_module=web,
error="rate_limit_exceeded",
include_ok=True,
)
allowed, err = require_admin_token(request)
+464 -44
View File
@@ -1,14 +1,83 @@
"""
Rate Limiting Service (S17).
Implements Token Bucket algorithm for per-IP rate limiting.
Rate Limiting Service (S17 / R143).
Provides shared request-scoped rate-limit evaluation with hierarchical budgets and
machine-readable diagnostics while preserving the legacy bool-only helper.
"""
from __future__ import annotations
import os
import threading
import time
from typing import Dict, Tuple
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Mapping, Optional, Tuple
from .request_ip import get_client_ip
try:
from .access_control import resolve_token_info
except ImportError:
from services.access_control import resolve_token_info # type: ignore
try:
from .tenant_context import DEFAULT_TENANT_ID, extract_tenant_from_headers
except ImportError:
from services.tenant_context import ( # type: ignore
DEFAULT_TENANT_ID,
extract_tenant_from_headers,
)
DEFAULT_RETRY_AFTER_SECONDS = 60
_REQUEST_CACHE_ATTR = "_openclaw_rate_limit_decisions"
_IP_SCALED_MULTIPLIER = 5.0
@dataclass(frozen=True)
class RateLimitDecision:
allowed: bool
limit_type: str
bucket: str
scope: str
retry_after_sec: int
reason_code: str
endpoint_class: str
ip: str
token_id: str = "anonymous"
tenant_id: str = DEFAULT_TENANT_ID
def to_payload(
self, *, error: str = "rate_limit_exceeded", include_ok: bool = True
) -> Dict[str, Any]:
payload: Dict[str, Any] = {
"error": error,
"code": "rate_limit_exceeded",
"bucket": self.bucket,
"scope": self.scope,
"retry_after_sec": self.retry_after_sec,
"reason_code": self.reason_code,
"endpoint_class": self.endpoint_class,
}
if include_ok:
payload["ok"] = False
return payload
@dataclass(frozen=True)
class BucketPolicy:
capacity: int
tokens_per_second: float
@dataclass(frozen=True)
class RateLimitPolicy:
principal: BucketPolicy
tenant: BucketPolicy
ip: BucketPolicy
endpoint_class: BucketPolicy
daily_cap_env: Optional[str] = None
class TokenBucket:
"""
@@ -18,89 +87,440 @@ class TokenBucket:
def __init__(self, capacity: int, tokens_per_second: float):
self.capacity = float(capacity)
self.tokens = float(capacity)
self.rate = tokens_per_second
self.rate = max(0.0, float(tokens_per_second))
self.last_update = time.time()
self.lock = threading.Lock()
def _refill_unlocked(self, now: float) -> None:
elapsed = max(0.0, now - self.last_update)
self.last_update = now
self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
def consume(self, amount: int = 1) -> bool:
"""
Attempt to consume tokens.
Returns True if successful, False if not enough tokens.
"""
allowed, _retry_after = self.consume_with_diagnostics(amount)
return allowed
def consume_with_diagnostics(self, amount: int = 1) -> Tuple[bool, int]:
"""
Attempt to consume tokens and return retry-after diagnostics on denial.
"""
with self.lock:
now = time.time()
elapsed = now - self.last_update
self.last_update = now
# Refill
self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
self._refill_unlocked(now)
if self.tokens >= amount:
self.tokens -= amount
return True
return False
return True, 0
if self.rate <= 0:
return False, DEFAULT_RETRY_AFTER_SECONDS
needed = amount - self.tokens
retry_after = int(max(1, (needed / self.rate) + 0.999999))
return False, retry_after
class DailyCounter:
"""UTC-day counter for optional daily caps."""
def __init__(self) -> None:
self._counts: Dict[str, Tuple[str, int]] = {}
self._lock = threading.Lock()
def check_and_increment(self, key: str, cap: int) -> Tuple[bool, int]:
if cap <= 0:
return True, 0
day_key = _utc_day_key()
with self._lock:
current_day, current_count = self._counts.get(key, (day_key, 0))
if current_day != day_key:
current_day, current_count = day_key, 0
if current_count >= cap:
return False, _seconds_until_next_utc_day()
self._counts[key] = (current_day, current_count + 1)
return True, 0
class RateLimiter:
"""
Manages rate limits for different endpoints/keys per client IP.
Manages hierarchical rate limits for different endpoint classes.
"""
def __init__(self):
# type -> IP -> Bucket
self.buckets: Dict[str, Dict[str, TokenBucket]] = {}
self.lock = threading.Lock()
# Default limits (capacity, rate/sec)
# rate = requests/minute / 60
self.daily_counters = DailyCounter()
self.policies = self._build_default_policies()
# IMPORTANT: preserve the legacy tuple map; older callers still inspect defaults directly.
self.defaults = {
"webhook": (30, 30.0 / 60.0), # 30 req/min
"logs": (60, 60.0 / 60.0), # 60 req/min
"admin": (20, 20.0 / 60.0), # 20 req/min
"bridge": (20, 20.0 / 60.0), # 20 req/min
# R101: New Quotas
"connector": (20, 20.0 / 60.0), # 20 req/min (aligned with bridge)
"trigger": (60, 60.0 / 60.0), # 60 req/min (higher due to automation)
limit_type: (
policy.principal.capacity,
policy.principal.tokens_per_second,
)
for limit_type, policy in self.policies.items()
}
def check(self, limit_type: str, ip: str) -> bool:
"""
Check if request is allowed for the given type and IP.
limit_type: "webhook", "logs", "admin"
"""
def _build_default_policies(self) -> Dict[str, RateLimitPolicy]:
# Base limits preserve the old default as the principal bucket. Tenant/IP and
# endpoint-class budgets widen above that so authenticated callers on a shared
# IP do not collide immediately on the legacy IP-only bucket.
def policy(
base_capacity: int,
*,
daily_env: Optional[str] = None,
) -> RateLimitPolicy:
base_rate = base_capacity / 60.0
return RateLimitPolicy(
principal=BucketPolicy(base_capacity, base_rate),
tenant=BucketPolicy(base_capacity * 3, base_rate * 3),
ip=BucketPolicy(base_capacity, base_rate),
endpoint_class=BucketPolicy(base_capacity * 10, base_rate * 10),
daily_cap_env=daily_env,
)
return {
"webhook": policy(30, daily_env="OPENCLAW_RATE_LIMIT_WEBHOOK_DAILY_CAP"),
"logs": policy(60),
"admin": policy(20, daily_env="OPENCLAW_RATE_LIMIT_ADMIN_DAILY_CAP"),
"bridge": policy(20, daily_env="OPENCLAW_RATE_LIMIT_BRIDGE_DAILY_CAP"),
"connector": policy(
20, daily_env="OPENCLAW_RATE_LIMIT_CONNECTOR_DAILY_CAP"
),
"trigger": policy(60, daily_env="OPENCLAW_RATE_LIMIT_TRIGGER_DAILY_CAP"),
"events": policy(30, daily_env="OPENCLAW_RATE_LIMIT_EVENTS_DAILY_CAP"),
}
def check(
self,
limit_type: str,
ip: str,
*,
token_id: Optional[str] = None,
tenant_id: Optional[str] = None,
) -> bool:
decision = self.evaluate(
limit_type,
ip,
token_id=token_id,
tenant_id=tenant_id,
)
return decision.allowed
def evaluate(
self,
limit_type: str,
ip: str,
*,
token_id: Optional[str] = None,
tenant_id: Optional[str] = None,
) -> RateLimitDecision:
ip = ip or "unknown"
token_id = (token_id or "").strip() or "anonymous"
tenant_id = (tenant_id or "").strip() or DEFAULT_TENANT_ID
policy = self.policies.get(limit_type, self.policies["webhook"])
# Determine config
capacity, rate = self.defaults.get(limit_type, (30, 0.5))
endpoint_decision = self._check_bucket(
limit_type,
bucket="endpoint_class",
scope_value=limit_type,
policy=policy.endpoint_class,
)
if not endpoint_decision.allowed:
return endpoint_decision
# Get or create bucket
# Use granular locking for bucket creation only
bucket = self._get_bucket(limit_type, ip, capacity, rate)
daily_cap = self._get_daily_cap(limit_type, policy.daily_cap_env)
if daily_cap:
principal_bucket, principal_scope = self._principal_scope(
token_id=token_id,
tenant_id=tenant_id,
ip=ip,
)
daily_allowed, retry_after = self.daily_counters.check_and_increment(
f"{limit_type}:{principal_bucket}:{principal_scope}",
daily_cap,
)
if not daily_allowed:
return RateLimitDecision(
allowed=False,
limit_type=limit_type,
bucket="daily",
scope=f"{principal_bucket}:{principal_scope}",
retry_after_sec=retry_after,
reason_code="daily_cap_exceeded",
endpoint_class=limit_type,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
return bucket.consume(1)
if token_id != "anonymous":
token_decision = self._check_bucket(
limit_type,
bucket="token_id",
scope_value=token_id,
policy=policy.principal,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
if not token_decision.allowed:
return token_decision
def _get_bucket(self, ltype: str, ip: str, cap: float, rate: float) -> TokenBucket:
# Double check locking pattern potentially but coarse lock is fine for dict access
if tenant_id != DEFAULT_TENANT_ID:
tenant_decision = self._check_bucket(
limit_type,
bucket="tenant",
scope_value=tenant_id,
policy=policy.tenant,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
if not tenant_decision.allowed:
return tenant_decision
ip_policy = policy.ip
if token_id != "anonymous" or tenant_id != DEFAULT_TENANT_ID:
ip_policy = BucketPolicy(
capacity=int(max(1, round(policy.ip.capacity * _IP_SCALED_MULTIPLIER))),
tokens_per_second=policy.ip.tokens_per_second * _IP_SCALED_MULTIPLIER,
)
ip_decision = self._check_bucket(
limit_type,
bucket="ip",
scope_value=ip,
policy=ip_policy,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
if not ip_decision.allowed:
return ip_decision
return RateLimitDecision(
allowed=True,
limit_type=limit_type,
bucket="allow",
scope=f"endpoint_class:{limit_type}",
retry_after_sec=0,
reason_code="allowed",
endpoint_class=limit_type,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
def _check_bucket(
self,
limit_type: str,
*,
bucket: str,
scope_value: str,
policy: BucketPolicy,
ip: str = "unknown",
token_id: str = "anonymous",
tenant_id: str = DEFAULT_TENANT_ID,
) -> RateLimitDecision:
bucket_obj = self._get_bucket(
limit_type,
bucket=bucket,
scope_value=scope_value,
capacity=policy.capacity,
rate=policy.tokens_per_second,
)
allowed, retry_after = bucket_obj.consume_with_diagnostics(1)
return RateLimitDecision(
allowed=allowed,
limit_type=limit_type,
bucket=bucket,
scope=f"{bucket}:{scope_value}",
retry_after_sec=retry_after if not allowed else 0,
reason_code="burst_limit_exceeded" if not allowed else "allowed",
endpoint_class=limit_type,
ip=ip,
token_id=token_id,
tenant_id=tenant_id,
)
def _get_bucket(
self,
limit_type: str,
*,
bucket: str,
scope_value: str,
capacity: int,
rate: float,
) -> TokenBucket:
bucket_type = f"{limit_type}:{bucket}"
with self.lock:
if ltype not in self.buckets:
self.buckets[ltype] = {}
typed = self.buckets.setdefault(bucket_type, {})
if scope_value not in typed:
typed[scope_value] = TokenBucket(capacity, rate)
return typed[scope_value]
if ip not in self.buckets[ltype]:
self.buckets[ltype][ip] = TokenBucket(cap, rate)
def _get_daily_cap(self, limit_type: str, env_name: Optional[str]) -> Optional[int]:
if not env_name:
return None
legacy_env = env_name.replace("OPENCLAW_", "MOLTBOT_", 1)
raw = (os.environ.get(env_name) or os.environ.get(legacy_env) or "").strip()
if not raw:
return None
try:
value = int(raw)
except ValueError:
return None
return value if value > 0 else None
return self.buckets[ltype][ip]
def _principal_scope(
self, *, token_id: str, tenant_id: str, ip: str
) -> Tuple[str, str]:
if token_id != "anonymous":
return "token_id", token_id
if tenant_id != DEFAULT_TENANT_ID:
return "tenant", tenant_id
return "ip", ip
def _utc_day_key(now: Optional[datetime] = None) -> str:
now = now or datetime.now(timezone.utc)
return now.strftime("%Y-%m-%d")
def _seconds_until_next_utc_day(now: Optional[datetime] = None) -> int:
now = now or datetime.now(timezone.utc)
tomorrow = (now + timedelta(days=1)).date()
next_day = datetime.combine(tomorrow, datetime.min.time(), tzinfo=timezone.utc)
delta = int((next_day - now).total_seconds())
return max(1, delta)
def resolve_rate_limit_context(request) -> Tuple[str, str, str]:
"""
Resolve stable request scope identifiers without leaking raw secrets.
"""
ip = get_client_ip(request) or "unknown"
token_id = "anonymous"
tenant_id = DEFAULT_TENANT_ID
try:
token_info = resolve_token_info(request)
except Exception:
token_info = None
if token_info is not None and getattr(token_info, "token_id", None):
token_id = str(getattr(token_info, "token_id") or "anonymous")
tenant_id = str(
getattr(token_info, "tenant_id", DEFAULT_TENANT_ID) or DEFAULT_TENANT_ID
)
headers = getattr(request, "headers", None)
if isinstance(headers, Mapping):
try:
header_tenant = extract_tenant_from_headers(headers)
except Exception:
header_tenant = None
if header_tenant:
tenant_id = header_tenant
return ip, token_id, tenant_id
# Global instance
rate_limiter = RateLimiter()
def _get_request_cache(request) -> Dict[str, RateLimitDecision]:
cache = getattr(request, _REQUEST_CACHE_ATTR, None)
if not isinstance(cache, dict):
cache = {}
try:
setattr(request, _REQUEST_CACHE_ATTR, cache)
except Exception:
return {}
return cache
def evaluate_rate_limit(request, limit_type: str) -> RateLimitDecision:
cache = _get_request_cache(request)
if limit_type in cache:
return cache[limit_type]
ip, token_id, tenant_id = resolve_rate_limit_context(request)
decision = rate_limiter.evaluate(
limit_type,
ip,
token_id=token_id,
tenant_id=tenant_id,
)
if cache is not None:
cache[limit_type] = decision
return decision
def get_cached_rate_limit_decision(
request, limit_type: str
) -> Optional[RateLimitDecision]:
cache = getattr(request, _REQUEST_CACHE_ATTR, None)
if isinstance(cache, dict):
decision = cache.get(limit_type)
if isinstance(decision, RateLimitDecision):
return decision
return None
def check_rate_limit(request, limit_type: str) -> bool:
"""
Helper to check rate limit from standard request object.
Returns True if allowed, False if exceeded.
"""
# S6: Resolve real IP
remote = get_client_ip(request)
return rate_limiter.check(limit_type, remote)
decision = evaluate_rate_limit(request, limit_type)
return decision.allowed
def build_rate_limit_payload(
request,
limit_type: str,
*,
error: str = "rate_limit_exceeded",
include_ok: bool = True,
) -> Dict[str, Any]:
decision = get_cached_rate_limit_decision(request, limit_type)
if decision is None:
decision = RateLimitDecision(
allowed=False,
limit_type=limit_type,
bucket="unknown",
scope=f"endpoint_class:{limit_type}",
retry_after_sec=DEFAULT_RETRY_AFTER_SECONDS,
reason_code="rate_limit_exceeded",
endpoint_class=limit_type,
ip="unknown",
)
return decision.to_payload(error=error, include_ok=include_ok)
def build_rate_limit_response(
request,
limit_type: str,
*,
web_module,
error: str = "rate_limit_exceeded",
include_ok: bool = True,
):
payload = build_rate_limit_payload(
request,
limit_type,
error=error,
include_ok=include_ok,
)
retry_after = str(payload.get("retry_after_sec", DEFAULT_RETRY_AFTER_SECONDS))
return web_module.json_response(
payload,
status=429,
headers={"Retry-After": retry_after},
)
+44 -3
View File
@@ -11,6 +11,7 @@ from services.failover import (
CooldownEntry,
ErrorCategory,
FailoverState,
classify_cooldown,
classify_error,
get_cooldown_duration,
get_failover_candidates,
@@ -55,6 +56,25 @@ class TestErrorClassification(unittest.TestCase):
ErrorCategory.RATE_LIMIT,
)
def test_retry_after_reason_code_is_preserved(self):
"""429 + retry-after should surface explicit cooldown diagnostics."""
class RetryAfterError(Exception):
retry_after = 42
decision = classify_cooldown(RetryAfterError("Too many requests"), 429)
self.assertEqual(decision.category, ErrorCategory.RATE_LIMIT)
self.assertEqual(decision.reason_code, "provider_retry_after")
self.assertEqual(decision.bucket, "provider_cooldown")
self.assertEqual(decision.retry_after_sec, 42)
def test_quota_reason_code_is_preserved(self):
"""Quota/billing style errors should map to provider_quota diagnostics."""
decision = classify_cooldown(Exception("Insufficient quota"), 429)
self.assertEqual(decision.category, ErrorCategory.BILLING)
self.assertEqual(decision.reason_code, "provider_quota_exceeded")
self.assertEqual(decision.bucket, "provider_quota")
def test_timeout_errors(self):
"""Should classify timeout errors."""
self.assertEqual(
@@ -145,7 +165,14 @@ class TestCooldownManagement(unittest.TestCase):
state = FailoverState(self.state_file)
# Set cooldown
state.set_cooldown("openai", "gpt-4", "rate_limit", 60)
state.set_cooldown(
"openai",
"gpt-4",
"rate_limit",
60,
reason_code="provider_rate_limited",
bucket="provider_cooldown",
)
# Should be in cooldown
self.assertTrue(state.is_cooling_down("openai", "gpt-4"))
@@ -156,7 +183,14 @@ class TestCooldownManagement(unittest.TestCase):
def test_cooldown_persistence(self):
"""Should persist cooldown state to disk."""
state1 = FailoverState(self.state_file)
state1.set_cooldown("openai", "gpt-4", "auth", 3600)
state1.set_cooldown(
"openai",
"gpt-4",
"auth",
3600,
reason_code="provider_auth_failed",
bucket="provider_auth",
)
# Create new instance (simulates restart)
state2 = FailoverState(self.state_file)
@@ -194,7 +228,14 @@ class TestCooldownManagement(unittest.TestCase):
def test_no_secrets_in_state(self):
"""Should not persist secrets in state file."""
state = FailoverState(self.state_file)
state.set_cooldown("openai", "gpt-4", "auth_failed", 60)
state.set_cooldown(
"openai",
"gpt-4",
"auth_failed",
60,
reason_code="provider_auth_failed",
bucket="provider_auth",
)
# Read raw file
with open(self.state_file, "r") as f:
+77 -1
View File
@@ -2,10 +2,17 @@
Tests for Rate Limiting Service (S17).
"""
import os
import time
import unittest
from unittest.mock import patch
from services.rate_limit import RateLimiter, TokenBucket
from services.rate_limit import (
RateLimitDecision,
RateLimiter,
TokenBucket,
build_rate_limit_payload,
)
class TestRateLimit(unittest.TestCase):
@@ -62,6 +69,75 @@ class TestRateLimit(unittest.TestCase):
# IP B should be fresh
self.assertTrue(limiter.check("webhook", "2.2.2.2"))
def test_rate_limiter_token_scope_isolates_shared_ip(self):
limiter = RateLimiter()
ip = "10.10.10.10"
for _ in range(20):
self.assertTrue(limiter.check("admin", ip, token_id="kid-a"))
self.assertFalse(limiter.check("admin", ip, token_id="kid-a"))
self.assertTrue(limiter.check("admin", ip, token_id="kid-b"))
def test_rate_limiter_tenant_scope_isolates_shared_ip(self):
limiter = RateLimiter()
ip = "10.10.10.20"
for _ in range(60):
self.assertTrue(limiter.evaluate("admin", ip, tenant_id="tenant-a").allowed)
self.assertFalse(limiter.evaluate("admin", ip, tenant_id="tenant-a").allowed)
self.assertTrue(limiter.evaluate("admin", ip, tenant_id="tenant-b").allowed)
def test_rate_limiter_daily_cap_returns_daily_reason(self):
limiter = RateLimiter()
with patch.dict(os.environ, {"OPENCLAW_RATE_LIMIT_ADMIN_DAILY_CAP": "2"}):
first = limiter.evaluate("admin", "9.9.9.9", token_id="kid-daily")
second = limiter.evaluate("admin", "9.9.9.9", token_id="kid-daily")
denied = limiter.evaluate("admin", "9.9.9.9", token_id="kid-daily")
self.assertTrue(first.allowed)
self.assertTrue(second.allowed)
self.assertFalse(denied.allowed)
self.assertEqual(denied.bucket, "daily")
self.assertEqual(denied.reason_code, "daily_cap_exceeded")
self.assertGreaterEqual(denied.retry_after_sec, 1)
def test_build_rate_limit_payload_includes_machine_readable_fields(self):
request = type("FakeRequest", (), {})()
setattr(
request,
"_openclaw_rate_limit_decisions",
{
"admin": RateLimitDecision(
allowed=False,
limit_type="admin",
bucket="token_id",
scope="token_id:kid-abc",
retry_after_sec=17,
reason_code="burst_limit_exceeded",
endpoint_class="admin",
ip="127.0.0.1",
token_id="kid-abc",
tenant_id="default",
)
},
)
payload = build_rate_limit_payload(
request,
"admin",
error="Rate limit exceeded",
include_ok=True,
)
self.assertFalse(payload["ok"])
self.assertEqual(payload["code"], "rate_limit_exceeded")
self.assertEqual(payload["bucket"], "token_id")
self.assertEqual(payload["scope"], "token_id:kid-abc")
self.assertEqual(payload["retry_after_sec"], 17)
self.assertEqual(payload["reason_code"], "burst_limit_exceeded")
if __name__ == "__main__":
unittest.main()