fix: rate limit 實施

Author: Mag1cFall

src/api/request_processor.py

@@ -292,6 +292,15 @@ async def create_stream_generator_from_helper(event_to_set: Event, task_to_cance
                                 continue
                         elif isinstance(raw_data, dict):
                             data = raw_data
+                            if data.get('error') == 'rate_limit':
+                                logger.warning(f"[{req_id}] 🚨 接收到来自代理的速率限制信号: {data}")
+                                try:
+                                    error_chunk = {'id': chat_completion_id, 'object': 'chat.completion.chunk', 'model': model_name_for_stream, 'created': created_timestamp, 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': f"\n\n[System: Rate Limit Exceeded - {data.get('detail', 'Quota exceeded')}]"}, 'finish_reason': 'stop', 'native_finish_reason': 'stop'}]}
+                                    yield f"data: {json.dumps(error_chunk, ensure_ascii=False, separators=(',', ':'))}\n\n"
+                                except: pass
+                                if not event_to_set.is_set():
+                                    event_to_set.set()
+                                break
                         else:
                             logger.warning(f'[{req_id}] 未知的流数据类型: {type(raw_data)}')
                             continue
@@ -339,6 +348,30 @@ async def create_stream_generator_from_helper(event_to_set: Event, task_to_cance
                                 choice_item = {'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': 'stop', 'native_finish_reason': 'stop'}
                             output = {'id': chat_completion_id, 'object': 'chat.completion.chunk', 'model': model_name_for_stream, 'created': created_timestamp, 'choices': [choice_item]}
                             yield f"data: {json.dumps(output, ensure_ascii=False, separators=(',', ':'))}\n\n"
+                    
+                    # Late Rate Limit Check
+                    late_check_wait = 2.0 if len(full_body_content) < 50 else 0.2
+                    if late_check_wait > 0.5:
+                         logger.info(f"[{req_id}] 内容较短 ({len(full_body_content)}), 等待 {late_check_wait}s 检查延迟 Rate Limit")
+                    await asyncio.sleep(late_check_wait)
+                    try:
+                        from server import STREAM_QUEUE
+                        import queue
+                        if STREAM_QUEUE:
+                            while True:
+                                try:
+                                    msg = STREAM_QUEUE.get_nowait()
+                                    if isinstance(msg, dict) and msg.get('error') == 'rate_limit':
+                                        logger.warning(f"[{req_id}] 🚨 捕获到延迟的 Rate Limit 信号: {msg}")
+                                        try:
+                                            error_chunk = {'id': chat_completion_id, 'object': 'chat.completion.chunk', 'model': model_name_for_stream, 'created': created_timestamp, 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': f"\n\n[System: Rate Limit Exceeded - {msg.get('detail', 'Quota exceeded')}]"}, 'finish_reason': 'stop', 'native_finish_reason': 'stop'}]}
+                                            yield f"data: {json.dumps(error_chunk, ensure_ascii=False, separators=(',', ':'))}\n\n"
+                                        except: pass
+                                except queue.Empty:
+                                    break
+                    except Exception as e:
+                        logger.error(f"[{req_id}] Late check failed: {e}")
+                    
                 except ClientDisconnectedError as disconnect_err:
                     abort_handler = AbortSignalHandler()
                     disconnect_info = abort_handler.handle_error(disconnect_err, req_id)
@@ -427,6 +460,9 @@ async def create_stream_generator_from_helper(event_to_set: Event, task_to_cance
                     continue
             elif isinstance(raw_data, dict):
                 data = raw_data
+                if data.get('error') == 'rate_limit':
+                    logger.warning(f"[{req_id}] 🚨 非流式请求中接收到速率限制: {data}")
+                    raise HTTPException(status_code=429, detail=f"Rate limit exceeded: {data.get('detail')}")
             else:
                 logger.warning(f'[{req_id}] 非流式未知数据类型: {type(raw_data)}')
                 continue

src/gateway.py

@@ -134,7 +134,7 @@ async def stream_proxy() -> AsyncGenerator[bytes, None]:
                         data, _ = chunk
                         if data:
                             check_count += 1
-                            if check_count <= 5 and not rate_limited:
+                            if not rate_limited:
                                 if check_rate_limit_in_response(data):
                                     rate_limited = True
                             yield data

src/proxy/server.py

@@ -45,6 +45,10 @@ def _get_tls_context(self, domain: str):
             certfile=self.cert_store.storage_dir / f'{domain}.crt',
             keyfile=self.cert_store.storage_dir / f'{domain}.key'
         )
+        try:
+            ctx.set_alpn_protocols(['http/1.1'])
+        except Exception:
+            pass
 
         if len(self._context_cache) > 50:
             self._context_cache.clear()
@@ -131,8 +135,13 @@ async def _process_tunnel(
             )
 
             try:
+                upstream_ctx = ssl.create_default_context()
+                try:
+                    upstream_ctx.set_alpn_protocols(['http/1.1'])
+                except Exception:
+                    pass
                 server_reader, server_writer = await self.connector.open_connection(
-                    host, port, ssl.create_default_context()
+                    host, port, upstream_ctx
                 )
                 await self._relay_with_inspection(
                     client_reader, client_writer,
@@ -219,7 +228,18 @@ async def process_upstream():
                             client_buf.clear()
                             continue
 
-                        if 'GenerateContent' in path:
+                        if 'jserror' in path:
+                            inspect_response = False
+                            try:
+                                path_str = path
+                                if 'quota' in path_str or 'limit' in path_str or 'exceeded' in path_str:
+                                    self.log.info(f"Rate limit keyword found in jserror: {path_str}")
+                                    if self.message_queue is not None:
+                                        self.message_queue.put({'error': 'rate_limit', 'detail': 'Rate limit detected via jserror', 'source': 'jserror', 'path': path_str})
+                            except Exception as e:
+                                self.log.error(f"Error inspecting jserror: {e}")
+                            server_writer.write(client_buf)
+                        elif 'GenerateContent' in path:
                             inspect_response = True
                             processed = await self.response_handler.handle_request(
                                 body_bytes, host, path