diff --git a/api/assist.py b/api/assist.py index 247e8ea..b7d9022 100644 --- a/api/assist.py +++ b/api/assist.py @@ -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() diff --git a/api/bridge.py b/api/bridge.py index a263778..a3b190b 100644 --- a/api/bridge.py +++ b/api/bridge.py @@ -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: diff --git a/api/checkpoints_handler.py b/api/checkpoints_handler.py index b93be0f..767c1ea 100644 --- a/api/checkpoints_handler.py +++ b/api/checkpoints_handler.py @@ -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: diff --git a/api/config.py b/api/config.py index 04b1049..779f7f8 100644 --- a/api/config.py +++ b/api/config.py @@ -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. diff --git a/api/connector_contracts.py b/api/connector_contracts.py index a0b4481..ea7235c 100644 --- a/api/connector_contracts.py +++ b/api/connector_contracts.py @@ -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: diff --git a/api/events.py b/api/events.py index 05a7af0..232eaac 100644 --- a/api/events.py +++ b/api/events.py @@ -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 diff --git a/api/preflight_handler.py b/api/preflight_handler.py index 4368fde..cb5c963 100644 --- a/api/preflight_handler.py +++ b/api/preflight_handler.py @@ -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) diff --git a/api/routes.py b/api/routes.py index 5efce57..57009ab 100644 --- a/api/routes.py +++ b/api/routes.py @@ -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: diff --git a/api/secrets.py b/api/secrets.py index 319e51d..9496ec8 100644 --- a/api/secrets.py +++ b/api/secrets.py @@ -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, ) diff --git a/api/security_doctor.py b/api/security_doctor.py index 1724bb5..b21b0a7 100644 --- a/api/security_doctor.py +++ b/api/security_doctor.py @@ -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 diff --git a/api/templates.py b/api/templates.py index 793904b..5511755 100644 --- a/api/templates.py +++ b/api/templates.py @@ -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: diff --git a/api/webhook.py b/api/webhook.py index 054417b..ed108ce 100644 --- a/api/webhook.py +++ b/api/webhook.py @@ -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: diff --git a/api/webhook_submit.py b/api/webhook_submit.py index ebe973a..6091c8b 100644 --- a/api/webhook_submit.py +++ b/api/webhook_submit.py @@ -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: diff --git a/api/webhook_validate.py b/api/webhook_validate.py index 0749fc2..33fe3bf 100644 --- a/api/webhook_validate.py +++ b/api/webhook_validate.py @@ -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 diff --git a/services/failover.py b/services/failover.py index 6ac45d4..4c1d011 100644 --- a/services/failover.py +++ b/services/failover.py @@ -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: diff --git a/services/llm_client.py b/services/llm_client.py index 217a227..13ade24 100644 --- a/services/llm_client.py +++ b/services/llm_client.py @@ -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, }, ) diff --git a/services/parameter_lab.py b/services/parameter_lab.py index bb98fed..a1690fc 100644 --- a/services/parameter_lab.py +++ b/services/parameter_lab.py @@ -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) diff --git a/services/rate_limit.py b/services/rate_limit.py index d09d3f0..007b291 100644 --- a/services/rate_limit.py +++ b/services/rate_limit.py @@ -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}, + ) diff --git a/tests/test_failover.py b/tests/test_failover.py index 84dc358..0c24c72 100644 --- a/tests/test_failover.py +++ b/tests/test_failover.py @@ -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: diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py index bb1614c..3764ed4 100644 --- a/tests/test_rate_limit.py +++ b/tests/test_rate_limit.py @@ -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()