修复首页限流及聊天请求额度混用

logistics_chat/hooks.py

@@ -7,11 +7,14 @@ from flask import current_app, jsonify, request
 
 def rate_limit():
     path = request.path
+    if request.endpoint == 'pages.index' and request.method in ('GET', 'HEAD'):
+        return None
     if path in ('/healthz', '/api/visitor/human/poll') or path.startswith('/static') or path.startswith('/api/human/'):
         return None
     ip = request.remote_addr or 'unknown'
     now = time.time()
-    hits = current_app.extensions['rate_hits'][ip]
+    bucket = (ip, 'chat') if request.endpoint == 'chat.stream_chat' else ip
+    hits = current_app.extensions['rate_hits'][bucket]
     while hits and now - hits[0] > 60:
         hits.pop(0)
     if len(hits) >= current_app.config['RATE_LIMIT_PER_MIN']:

tests/test_api.py

@@ -359,12 +359,44 @@ class ApiTests(unittest.TestCase):
 
     def test_health_check_remains_available_when_business_requests_are_rate_limited(self):
         self.app.config['RATE_LIMIT_PER_MIN'] = 1
-        self.assertEqual(self.client.get('/').status_code, 200)
+        self.assertEqual(self.client.get('/api/get_history').status_code, 200)
         self.assertEqual(self.client.get('/api/get_history').status_code, 429)
         response = self.client.get('/healthz')
         self.assertEqual(response.status_code, 200)
         self.assertEqual(response.json, {'status': 'ok'})
 
+    def test_homepage_does_not_consume_or_depend_on_business_rate_budget(self):
+        self.app.config['RATE_LIMIT_PER_MIN'] = 1
+        for method in ('get', 'head'):
+            response = getattr(self.client, method)('/')
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(response.mimetype, 'text/html')
+        self.assertEqual(self.client.get('/api/get_history').status_code, 200)
+        self.assertEqual(self.client.get('/api/get_history').status_code, 429)
+        for method in ('get', 'head'):
+            response = getattr(self.client, method)('/')
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(response.mimetype, 'text/html')
+
+    def test_chat_and_supporting_requests_have_separate_limited_budgets(self):
+        self.app.config['RATE_LIMIT_PER_MIN'] = 1
+        response = self.client.post('/stream_chat', json=[])
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.mimetype, 'text/event-stream')
+        self.assertEqual(self.client.post('/stream_chat', json=[]).status_code, 429)
+        self.assertEqual(self.client.get('/api/get_history').status_code, 200)
+        self.assertEqual(self.client.get('/api/get_history').status_code, 429)
+        self.assertEqual(self.client.get('/stream_chat').status_code, 429)
+
+    def test_supporting_request_limit_does_not_block_chat(self):
+        self.app.config['RATE_LIMIT_PER_MIN'] = 1
+        self.assertEqual(self.client.get('/api/get_history').status_code, 200)
+        self.assertEqual(self.client.get('/api/get_history').status_code, 429)
+        response = self.client.post('/stream_chat', json=[])
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.mimetype, 'text/event-stream')
+        self.assertEqual(self.client.post('/stream_chat', json=[]).status_code, 429)
+
     def test_feedback_stats_and_csv(self):
         cid = self.conversation()
         with self.app.app_context():