From b1555e1359bfd42aaa55a617abc1d8614a778e0d Mon Sep 17 00:00:00 2001 From: TheAchiever6823 <61914223+ShawnMadadha@users.noreply.github.com> Date: Sun, 17 May 2026 19:40:12 -0700 Subject: [PATCH] [shawn] feat: OpenSwarm-side rate limiting for LinkedIn MCP via pre_tool_hook --- backend/apps/agents/agent_manager.py | 22 +++ backend/apps/tools_lib/mcp_rate_limiter.py | 153 +++++++++++++++++++++ 2 files changed, 175 insertions(+) create mode 100644 backend/apps/tools_lib/mcp_rate_limiter.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index fde3380b..dd0f28dc 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -24,6 +24,7 @@ from backend.apps.tools_lib.tools_lib import ( refresh_google_token, refresh_hubspot_token, ) +from backend.apps.tools_lib import mcp_rate_limiter from backend.config.paths import SESSIONS_DIR from backend.apps.service.client import sync as _sync @@ -1282,6 +1283,27 @@ class AgentManager: tool_name = input_data.get("tool_name", "") hook_event = input_data.get("hook_event_name", "PreToolUse") + # OpenSwarm-side proactive rate limiting for MCP tools. Sits in + # front of the per-tool permission check so caps apply even to + # tools the user has set to "always_allow". Servers with their + # own rigorous limiter (e.g. vendored Instagram) are not in + # POLICIES, so this is a no-op for them. + if tool_name.startswith("mcp__"): + _m = re.match(r"mcp__([^_]+(?:-[^_]+)*)__(.+)", tool_name) + if _m: + _mcp_server, _mcp_tool = _m.group(1), _m.group(2) + _rl = mcp_rate_limiter.check(_mcp_server, _mcp_tool) + if _rl and "deny" in _rl: + return { + "hookSpecificOutput": { + "hookEventName": hook_event, + "permissionDecision": "deny", + "permissionDecisionReason": _rl["deny"], + } + } + if _rl and _rl.get("jitter_s", 0) > 0: + await asyncio.sleep(_rl["jitter_s"]) + if tool_name and tool_name != "AskUserQuestion": tool_input = input_data.get("tool_input", {}) policy = _maybe_override_policy( diff --git a/backend/apps/tools_lib/mcp_rate_limiter.py b/backend/apps/tools_lib/mcp_rate_limiter.py new file mode 100644 index 00000000..3d9244ac --- /dev/null +++ b/backend/apps/tools_lib/mcp_rate_limiter.py @@ -0,0 +1,153 @@ +"""OpenSwarm-side rate limiting for MCP tool calls. + +Sits in the pre_tool_hook so we can refuse a call before the SDK proxies it +to the MCP server. Per-server, per-category caps + randomized jitter + +persistent state across OpenSwarm restarts. + +Why this exists separately from per-server limiters: the LinkedIn upstream +(stickerdaniel/linkedin-mcp-server) only does reactive rate-limit detection +once LinkedIn already blocked us. We want proactive caps that stop us from +getting to that point. Same defense model as the vendored Instagram +server's built-in limiter, but at OpenSwarm's dispatch layer so we can +protect any third-party MCP server without forking it. + +Servers with their own rigorous limiter (currently: Instagram) are +deliberately not in POLICIES here, to avoid double-throttling. +""" +from __future__ import annotations + +import fnmatch +import json +import logging +import os +import random +import time +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +_STATE_PATH = Path.home() / ".openswarm" / "mcp-rate-limits.json" + +# Per-MCP-server policy. Keys match the sanitized server name from the +# `mcp____` SDK format. Each tool pattern is glob-matched +# top-to-bottom; first match wins, "*" is the fallback. +POLICIES: dict[str, dict] = { + "linkedin": { + "categories": { + "dm_send": {"per_minute": 2, "per_hour": 20, "per_day": 60, "jitter": (1.5, 4.0)}, + "connect": {"per_minute": 1, "per_hour": 10, "per_day": 40, "jitter": (2.0, 5.0)}, + "search": {"per_minute": 10, "per_hour": 60, "per_day": 300, "jitter": (0.5, 2.0)}, + "lookup": {"per_minute": 15, "per_hour": 120, "per_day": 600, "jitter": (0.3, 1.0)}, + }, + "tools": [ + ("send_message", "dm_send"), + ("connect_with_person", "connect"), + ("search_*", "search"), + ("*", "lookup"), + ], + }, +} + + +def _env_override(server: str, category: str, key: str, default: int) -> int: + var = f"{server.upper()}_RATE_LIMIT_{category.upper()}_{key.upper()}" + raw = os.environ.get(var) + if not raw: + return default + try: + value = int(raw) + return value if value > 0 else default + except ValueError: + return default + + +def _resolve(server: str, tool: str) -> tuple[str, dict] | None: + policy = POLICIES.get(server.lower()) + if not policy: + return None + for pattern, cat in policy["tools"]: + if fnmatch.fnmatch(tool, pattern): + base = policy["categories"].get(cat) + if not base: + return None + limits = { + "per_minute": _env_override(server, cat, "per_minute", base["per_minute"]), + "per_hour": _env_override(server, cat, "per_hour", base["per_hour"]), + "per_day": _env_override(server, cat, "per_day", base["per_day"]), + "jitter": base["jitter"], + } + return (cat, limits) + return None + + +def _load_state() -> dict[str, list[float]]: + if not _STATE_PATH.exists(): + return {} + try: + data = json.loads(_STATE_PATH.read_text()) + return {k: [float(t) for t in v] for k, v in data.items() if isinstance(v, list)} + except Exception: + return {} + + +def _save_state(state: dict[str, list[float]]) -> None: + try: + _STATE_PATH.parent.mkdir(parents=True, exist_ok=True) + _STATE_PATH.write_text(json.dumps(state)) + except Exception as exc: + logger.warning(f"Could not persist MCP rate-limit state: {exc}") + + +def _prune(timestamps: list[float], now: float, window_s: int) -> list[float]: + cutoff = now - window_s + return [t for t in timestamps if t >= cutoff] + + +def _fmt_duration(seconds: int) -> str: + if seconds < 60: + return f"{seconds}s" + if seconds < 3600: + return f"{seconds // 60}m {seconds % 60}s" + return f"{seconds // 3600}h {(seconds % 3600) // 60}m" + + +def check(mcp_server: str, tool: str) -> dict[str, Any] | None: + """Decide whether an MCP tool call should be allowed. + + Returns: + None — no policy for this server (skip rate limiting entirely) + {"allow": True, "jitter_s": float} — allow; caller should sleep that long + {"deny": str, "retry_after_seconds": int} — block; caller should refuse + """ + resolved = _resolve(mcp_server, tool) + if not resolved: + return None + category, limits = resolved + state_key = f"{mcp_server}:{category}" + + state = _load_state() + now = time.time() + stamps = _prune(state.get(state_key, []), now, 24 * 3600) + + for window_name, window_s in (("per_minute", 60), ("per_hour", 3600), ("per_day", 86400)): + in_window = _prune(stamps, now, window_s) + cap = limits[window_name] + if len(in_window) >= cap: + oldest = min(in_window) + retry_after = int((oldest + window_s) - now) + 1 + label = window_name.replace("per_", "") + reason = ( + f"OpenSwarm rate limit ({mcp_server}/{category}): {cap}/{label} cap reached " + f"(currently {len(in_window)}). Retry in {_fmt_duration(retry_after)}. " + f"This protects the connected account from anti-abuse bans. Override with " + f"{mcp_server.upper()}_RATE_LIMIT_{category.upper()}_{window_name.upper()}=N." + ) + logger.warning(f"[mcp-rate-limit] BLOCK {mcp_server}/{tool}: {reason}") + return {"deny": reason, "retry_after_seconds": retry_after} + + state[state_key] = stamps + [now] + _save_state(state) + lo, hi = limits["jitter"] + jitter = random.uniform(lo, hi) if hi > 0 else 0.0 + return {"allow": True, "jitter_s": jitter, "category": category}