mcp 限流优化

public_server.py

@@ -28,10 +28,17 @@ class PublicMcpHttpHandler:
         # unknown names fall back to the bare IP bucket, preventing bucket explosion
         self._known_tools = frozenset(t.get('name', '') for t in gateway_app.list_tools())
 
-    def _check_rate_limit(self, client_ip, rate_key, method, request_id=None):
-        """Returns an error response if rate limit exceeded, else None."""
-        if self.rate_limiter and client_ip and not self.rate_limiter.is_allowed(rate_key):
-            logger.warning(f"[RATE_LIMIT] rate limit exceeded: ip={client_ip}, method={method}")
+    def _check_rate_limit(self, rate_key, method, request_id=None, *, log_identity=None):
+        """Returns an error response if rate limit exceeded, else None.
+
+        rate_key     — the bucket key used by the limiter (session-based for tools/call,
+                       IP-based for tools/list)
+        log_identity — optional string shown in warning logs (e.g. client_ip)
+        """
+        if self.rate_limiter and rate_key and not self.rate_limiter.is_allowed(rate_key):
+            logger.warning(
+                f"[RATE_LIMIT] rate limit exceeded: identity={log_identity or rate_key}, method={method}"
+            )
             return McpProtocolHandler._error_response(request_id, -32000, 'Rate limit exceeded. Please try again later.')
         return None
 
@@ -52,8 +59,12 @@ class PublicMcpHttpHandler:
                     },
                 })
             if method == 'tools/list':
-                # Use a dedicated bucket for list so it doesn't compete with tools/call quota
-                blocked = self._check_rate_limit(client_ip, 'list:{0}'.format(client_ip), method, request_id)
+                # tools/list has no session context — use a dedicated IP-based bucket
+                # so it doesn't compete with the per-session tools/call quota
+                blocked = self._check_rate_limit(
+                    'list:{0}'.format(client_ip), method, request_id,
+                    log_identity=client_ip,
+                )
                 if blocked:
                     return blocked
                 tools = []
@@ -64,19 +75,21 @@ class PublicMcpHttpHandler:
                     tools.append(normalized)
                 return McpProtocolHandler._success_response(request_id, {'tools': tools})
             if method == 'tools/call':
-                # Per-tool independent quota: only registered tool names get their own bucket;
-                # unrecognized names fall into the shared IP bucket to prevent bucket explosion
-                tool_name = ((message.get('params') or {}).get('name') or '').strip()
-                rate_key = '{0}:{1}'.format(client_ip, tool_name) if tool_name in self._known_tools else client_ip
-                blocked = self._check_rate_limit(client_ip, rate_key, method, request_id)
-                if blocked:
-                    return blocked
+                # Parse context first — tools/call always requires a valid session
                 context = self.context_parser.parse(headers or {})
                 if not context.has_session():
                     raise RuntimeError(DEVICE_INVALID_MESSAGE)
+                # Session-based rate limiting: each employee gets an independent quota per tool.
+                # Unknown tool names fall back to the bare session bucket to prevent key explosion.
+                session_id = context.gateway_session_id
+                tool_name = ((message.get('params') or {}).get('name') or '').strip()
+                rate_key = '{0}:{1}'.format(session_id, tool_name) if tool_name in self._known_tools else session_id
+                blocked = self._check_rate_limit(rate_key, method, request_id, log_identity=client_ip)
+                if blocked:
+                    return blocked
                 params = message.get('params') or {}
                 result = self.gateway_app.call_tool(
-                    context.gateway_session_id,
+                    session_id,
                     params.get('name'),
                     params.get('arguments') or {},
                     request_id='rq_http_{0}_{1}'.format(request_id, uuid.uuid4().hex[:16]),
@@ -161,7 +174,7 @@ def serve_public(gateway_app, host='0.0.0.0', port=8765, enable_rate_limit=True,
             max_requests=rate_limit_max_requests,
             window_seconds=rate_limit_window_seconds,
         )
-        logger.info(f"Rate limiting enabled: {rate_limit_max_requests} requests/{rate_limit_window_seconds}s per IP")
+        logger.info(f"Rate limiting enabled: {rate_limit_max_requests} requests/{rate_limit_window_seconds}s per session (tools/call) / per IP (tools/list)")
 
     logger.info(f"Starting public MCP Gateway on {host}:{port}")
     server = ThreadingHTTPServer((host, int(port)), create_http_handler(gateway_app, rate_limiter))

tests/test_public_server.py

@@ -107,22 +107,33 @@ class RateLimitTest(unittest.TestCase):
         response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('query_track'), client_ip='1.2.3.4')
         self.assertNotIn('error', response)
 
-    def test_unknown_tool_name_uses_shared_ip_bucket_not_new_bucket(self):
-        # 未知工具名应归入 client_ip 桶，不能无限新建桶
+    def test_unknown_tool_name_uses_shared_session_bucket_not_new_bucket(self):
+        # 未知工具名应归入 session_id 桶，同一会话的不同未知工具共享同一配额，不能无限新建桶
         handler = self._make_handler(max_requests=1)
-        # 先用 client_ip 桶打一次（用未知工具名触发）
-        handler.handle_json_rpc({}, self._tools_call_msg('fake_tool_1'), client_ip='1.2.3.4')
-        # 再用另一个未知工具名，应该命中同一个IP桶，被限流
-        response = handler.handle_json_rpc({}, self._tools_call_msg('fake_tool_2'), client_ip='1.2.3.4')
+        # 先用 session 桶打一次（用未知工具名触发）
+        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_1'), client_ip='1.2.3.4')
+        # 再用同会话的另一个未知工具名，应该命中同一个 session 桶，被限流
+        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_2'), client_ip='1.2.3.4')
         self.assertIn('error', response)
         self.assertIn('Rate limit', response['error']['message'])
 
-    def test_different_ips_have_independent_quotas(self):
+    def test_different_sessions_have_independent_quotas(self):
+        # 不同会话（不同员工）即使来自同一IP，也拥有独立的配额
         handler = self._make_handler(max_requests=1)
         handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1')
-        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='2.2.2.2')
+        # 不同 session，即使同一 IP，quota 独立 → 应该允许
+        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_B'}, self._tools_call_msg(), client_ip='1.1.1.1')
         self.assertNotIn('error', response)
 
+    def test_same_session_from_different_ips_shares_quota(self):
+        # 同一会话从不同IP发来（如移动网络切换），仍共享同一会话配额
+        handler = self._make_handler(max_requests=1)
+        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1')
+        # 同 session，不同 IP → 命中同一 session 桶，被限流
+        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='2.2.2.2')
+        self.assertIn('error', response)
+        self.assertIn('Rate limit', response['error']['message'])
+
     # --- initialize 不受限流，tools/list 有独立 bucket ---
 
     def test_initialize_is_never_rate_limited(self):