refactor: modular architecture v2 with grid logging system and comprehensive test suite

Author: NikkeTryHard

README.md

@@ -561,14 +561,12 @@ python gui_launcher.py
 #### 响应获取优先级
 
 1. **第一层: 集成流式代理服务** (默认启用)
-
    - ⚡ **性能最优**: 直接转发流式响应，延迟最低
    - 📍 **端口**: 3120 (可通过 `STREAM_PORT` 配置)
    - ✅ **适用场景**: 流式请求、实时对话
    - ⚠️ **限制**: 参数支持有限，主要支持基础参数
 
 2. **第二层: 外部 Helper 服务** (可选配置)
-
    - 🔧 **需要**: 有效的认证文件 (`auth_profiles/active/*.json`)
    - 📡 **配置**: 通过 `--helper <endpoint>` 或 `.env` 配置
    - ✅ **适用场景**: 需要额外功能的场景

api_utils/app.py

@@ -17,6 +17,9 @@
 
 import stream
 
+# --- 导入集中状态模块 ---
+from api_utils.server_state import state
+
 # --- browser_utils模块导入 ---
 from browser_utils import (
     _close_page_logic,  # pyright: ignore[reportPrivateUsage]
@@ -35,9 +38,6 @@
 # --- models模块导入 ---
 from models import WebSocketConnectionManager
 
-# --- 导入集中状态模块 ---
-from api_utils.server_state import state
-
 from . import auth_utils
 
 
@@ -167,17 +167,23 @@ async def _shutdown_resources() -> None:
     logger = state.logger
     logger.info("Shutting down resources...")
 
+    # Signal all streaming generators to exit immediately
+    state.should_exit = True
+
     if state.STREAM_PROCESS:
         state.STREAM_PROCESS.terminate()
         logger.info("STREAM proxy terminated.")
 
     if state.worker_task and not state.worker_task.done():
+        logger.info("Cancelling worker task...")
         state.worker_task.cancel()
         try:
-            await asyncio.wait_for(state.worker_task, timeout=5.0)
-        except (asyncio.TimeoutError, asyncio.CancelledError):
-            pass
-        logger.info("Worker task stopped.")
+            await asyncio.wait_for(state.worker_task, timeout=2.0)
+            logger.info("Worker task cancelled.")
+        except asyncio.TimeoutError:
+            logger.warning("Worker task did not respond to cancellation within 2s.")
+        except asyncio.CancelledError:
+            logger.info("Worker task cancelled.")
 
     if state.page_instance:
         await _close_page_logic()

api_utils/client_connection.py

@@ -5,6 +5,7 @@
 
 from fastapi import HTTPException, Request
 
+from logging_utils import set_request_id
 from models import ClientDisconnectedError
 
 
@@ -61,14 +62,15 @@ async def setup_disconnect_monitoring(
 
     logger = logging.getLogger("AIStudioProxyServer")
     client_disconnected_event = Event()
+    set_request_id(req_id)
 
     async def check_disconnect_periodically() -> None:
         while not client_disconnected_event.is_set():
             try:
                 is_connected = await check_client_connection(req_id, http_request)
                 if not is_connected:
                     logger.info(
-                        f"[{req_id}] Active disconnect check detected client disconnection."
+                        "Active disconnect check detected client disconnection."
                     )
                     client_disconnected_event.set()
                     if not result_future.done():
@@ -85,7 +87,7 @@ async def check_disconnect_periodically() -> None:
                 # Task cancelled, exit gracefully
                 break
             except Exception as e:
-                logger.error(f"[{req_id}] (Disco Check Task) Error: {e}")
+                logger.error(f"(Disco Check Task) Error: {e}")
                 client_disconnected_event.set()
                 if not result_future.done():
                     result_future.set_exception(
@@ -100,7 +102,7 @@ async def check_disconnect_periodically() -> None:
 
     def check_client_disconnected(stage: str = "") -> bool:
         if client_disconnected_event.is_set():
-            logger.info(f"[{req_id}] Client disconnected detected at stage: '{stage}'")
+            logger.info(f"Client disconnected detected at stage: '{stage}'")
             raise ClientDisconnectedError(
                 f"[{req_id}] Client disconnected at stage: {stage}"
             )
@@ -116,13 +118,14 @@ async def enhanced_disconnect_monitor(
     Enhanced disconnect monitor for streaming responses.
     Returns True if client disconnected early.
     """
+    set_request_id(req_id)
     client_disconnected_early = False
     while not completion_event.is_set():
         try:
             is_connected = await check_client_connection(req_id, http_request)
             if not is_connected:
                 logger.info(
-                    f"[{req_id}] (Monitor) Client disconnected during streaming, triggering completion event."
+                    "(Monitor) Client disconnected during streaming, triggering completion event."
                 )
                 client_disconnected_early = True
                 if not completion_event.is_set():
@@ -132,7 +135,7 @@ async def enhanced_disconnect_monitor(
         except asyncio.CancelledError:
             break
         except Exception as e:
-            logger.error(f"[{req_id}] (Monitor) Enhanced disconnect checker error: {e}")
+            logger.error(f"(Monitor) Enhanced disconnect checker error: {e}")
             break
     return client_disconnected_early
 
@@ -147,13 +150,14 @@ async def non_streaming_disconnect_monitor(
     Disconnect monitor for non-streaming responses.
     Returns True if client disconnected early.
     """
+    set_request_id(req_id)
     client_disconnected_early = False
     while not result_future.done():
         try:
             is_connected = await check_client_connection(req_id, http_request)
             if not is_connected:
                 logger.info(
-                    f"[{req_id}] (Monitor) Client disconnected during non-streaming processing."
+                    "(Monitor) Client disconnected during non-streaming processing."
                 )
                 client_disconnected_early = True
                 if not result_future.done():
@@ -168,8 +172,6 @@ async def non_streaming_disconnect_monitor(
         except asyncio.CancelledError:
             break
         except Exception as e:
-            logger.error(
-                f"[{req_id}] (Monitor) Non-streaming disconnect checker error: {e}"
-            )
+            logger.error(f"(Monitor) Non-streaming disconnect checker error: {e}")
             break
     return client_disconnected_early

api_utils/context_init.py

@@ -1,3 +1,4 @@
+from logging_utils import set_request_id
 from models import ChatCompletionRequest
 
 from .context_types import RequestContext
@@ -6,31 +7,21 @@
 async def initialize_request_context(
     req_id: str, request: ChatCompletionRequest
 ) -> RequestContext:
-    from server import (
-        current_ai_studio_model_id,
-        is_page_ready,
-        logger,
-        model_switching_lock,
-        page_instance,
-        page_params_cache,
-        params_cache_lock,
-        parsed_model_list,
-    )
+    from api_utils.server_state import state
 
-    logger.info(f"[{req_id}] 开始处理请求...")
-    logger.info(
-        f"[{req_id}]   请求参数 - Model: {request.model}, Stream: {request.stream}"
-    )
+    set_request_id(req_id)
+    state.logger.info("开始处理请求...")
+    state.logger.info(f"  请求参数 - Model: {request.model}, Stream: {request.stream}")
 
     context: RequestContext = {
-        "logger": logger,
-        "page": page_instance,
-        "is_page_ready": is_page_ready,
-        "parsed_model_list": parsed_model_list,
-        "current_ai_studio_model_id": current_ai_studio_model_id,
-        "model_switching_lock": model_switching_lock,
-        "page_params_cache": page_params_cache,
-        "params_cache_lock": params_cache_lock,
+        "logger": state.logger,
+        "page": state.page_instance,
+        "is_page_ready": state.is_page_ready,
+        "parsed_model_list": state.parsed_model_list,
+        "current_ai_studio_model_id": state.current_ai_studio_model_id,
+        "model_switching_lock": state.model_switching_lock,
+        "page_params_cache": state.page_params_cache,
+        "params_cache_lock": state.params_cache_lock,
         "is_streaming": request.stream,
         "model_actually_switched": False,
         "requested_model": request.model,

api_utils/model_switching.py

@@ -1,20 +1,22 @@
 from playwright.async_api import Page as AsyncPage
 
 from api_utils.server_state import state
+from logging_utils import set_request_id
 
 from .context_types import RequestContext
 
 
 async def analyze_model_requirements(
     req_id: str, context: RequestContext, requested_model: str, proxy_model_name: str
 ) -> RequestContext:
+    set_request_id(req_id)
     logger = context["logger"]
     current_ai_studio_model_id = context["current_ai_studio_model_id"]
     parsed_model_list = context["parsed_model_list"]
 
     if requested_model and requested_model != proxy_model_name:
         requested_model_id = requested_model.split("/")[-1]
-        logger.info(f"[{req_id}] 请求使用模型: {requested_model_id}")
+        logger.info(f"请求使用模型: {requested_model_id}")
 
         if parsed_model_list:
             valid_model_ids = [m.get("id") for m in parsed_model_list]
@@ -30,7 +32,7 @@ async def analyze_model_requirements(
         if current_ai_studio_model_id != requested_model_id:
             context["needs_model_switching"] = True
             logger.info(
-                f"[{req_id}] 需要切换模型: 当前={current_ai_studio_model_id} -> 目标={requested_model_id}"
+                f"需要切换模型: 当前={current_ai_studio_model_id} -> 目标={requested_model_id}"
             )
 
     return context
@@ -39,6 +41,7 @@ async def analyze_model_requirements(
 async def handle_model_switching(
     req_id: str, context: RequestContext
 ) -> RequestContext:
+    set_request_id(req_id)
     if not context["needs_model_switching"]:
         return context
 
@@ -50,7 +53,7 @@ async def handle_model_switching(
     async with model_switching_lock:
         if state.current_ai_studio_model_id != model_id_to_use:
             logger.info(
-                f"[{req_id}] 准备切换模型: {state.current_ai_studio_model_id} -> {model_id_to_use}"
+                f"准备切换模型: {state.current_ai_studio_model_id} -> {model_id_to_use}"
             )
             from browser_utils import switch_ai_studio_model
 
@@ -59,9 +62,7 @@ async def handle_model_switching(
                 state.current_ai_studio_model_id = model_id_to_use
                 context["model_actually_switched"] = True
                 context["current_ai_studio_model_id"] = model_id_to_use
-                logger.info(
-                    f"[{req_id}] 模型切换成功: {state.current_ai_studio_model_id}"
-                )
+                logger.info(f"模型切换成功: {state.current_ai_studio_model_id}")
             else:
                 await _handle_model_switch_failure(
                     req_id,
@@ -77,7 +78,8 @@ async def handle_model_switching(
 async def _handle_model_switch_failure(
     req_id: str, page: AsyncPage, model_id_to_use: str, model_before_switch: str, logger
 ) -> None:
-    logger.warning(f"[{req_id}] 模型切换至 {model_id_to_use} 失败。")
+    set_request_id(req_id)
+    logger.warning(f"模型切换至 {model_id_to_use} 失败。")
     state.current_ai_studio_model_id = model_before_switch
     from .error_utils import http_error
 
@@ -87,6 +89,7 @@ async def _handle_model_switch_failure(
 
 
 async def handle_parameter_cache(req_id: str, context: RequestContext) -> None:
+    set_request_id(req_id)
     logger = context["logger"]
     params_cache_lock = context["params_cache_lock"]
     page_params_cache = context["page_params_cache"]
@@ -100,7 +103,7 @@ async def handle_parameter_cache(req_id: str, context: RequestContext) -> None:
         if model_actually_switched or (
             current_ai_studio_model_id != cached_model_for_params
         ):
-            logger.info(f"[{req_id}] 模型已更改，参数缓存失效。")
+            logger.info("模型已更改，参数缓存失效。")
             page_params_cache.clear()
             page_params_cache["last_known_model_id_for_params"] = (
                 current_ai_studio_model_id

api_utils/page_response.py

@@ -5,21 +5,23 @@
 from playwright.async_api import expect as expect_async
 
 from config import RESPONSE_CONTAINER_SELECTOR, RESPONSE_TEXT_SELECTOR
+from logging_utils import set_request_id
 
 
 async def locate_response_elements(
     page: AsyncPage, req_id: str, logger, check_client_disconnected
 ) -> None:
     """定位响应容器与文本元素，包含超时与错误处理。"""
-    logger.info(f"[{req_id}] 定位响应元素...")
+    set_request_id(req_id)
+    logger.info("定位响应元素...")
     response_container = page.locator(RESPONSE_CONTAINER_SELECTOR).last
     response_element = response_container.locator(RESPONSE_TEXT_SELECTOR)
 
     try:
         await expect_async(response_container).to_be_attached(timeout=20000)
         check_client_disconnected("After Response Container Attached: ")
         await expect_async(response_element).to_be_attached(timeout=90000)
-        logger.info(f"[{req_id}] 响应元素已定位。")
+        logger.info("响应元素已定位。")
     except (PlaywrightAsyncError, asyncio.TimeoutError) as locate_err:
         from .error_utils import upstream_error