mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 00:48:07 +00:00
refactor: harden rate limit diagnostics and cooldowns
This commit is contained in:
+16
-4
@@ -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
@@ -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:
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user