Merge pull request #271 from NikkeTryHard/feature/monkeytype-strict-typing

refactor: comprehensive strict type annotations with enhanced test coverage

Author: NikkeTryHard

.gitignore

@@ -188,6 +188,7 @@ coverage.xml
 *.py,cover
 .hypothesis/
 .pytest_cache/
+.testmondata*
 
 # Environments
 .env
@@ -260,3 +261,12 @@ browser_utils/generated_*.js
 # Docker 环境的实际配置文件（保留示例文件）
 docker/.env
 docker/my_*.json
+
+monkeytype.sqlite3
+
+# Temporary debug/output files
+pyright_output.txt
+pyright_utils_full.txt
+temp_*.txt
+temp_*.md
+utils_errors.txt

api_utils/app.py

@@ -228,6 +228,8 @@ async def lifespan(app: FastAPI):
         logger.info("Server startup complete.")
         state.is_initializing = False
         yield
+    except asyncio.CancelledError:
+        raise
     except Exception as e:
         logger.critical(f"Application startup failed: {e}", exc_info=True)
         await _shutdown_resources()

api_utils/client_connection.py

@@ -1,14 +1,22 @@
 import asyncio
 from asyncio import Event, Task
 from logging import Logger
-from typing import Any, Callable, Coroutine, Tuple
+from typing import Any, Callable, Coroutine, Dict, Protocol, Tuple
 
 from fastapi import HTTPException, Request
 
 from logging_utils import set_request_id
 from models import ClientDisconnectedError
 
 
+class SupportsReceive(Protocol):
+    """Protocol for request objects that support _receive method."""
+
+    def _receive(self) -> Coroutine[Any, Any, Dict[str, Any]]:
+        """Internal method to receive messages from ASGI."""
+        ...
+
+
 async def check_client_connection(req_id: str, http_request: Request) -> bool:
     """
     Checks if the client is still connected.
@@ -18,9 +26,12 @@ async def check_client_connection(req_id: str, http_request: Request) -> bool:
         if hasattr(http_request, "_receive"):
             try:
                 # Use a very short timeout to check for disconnect message
-                # Cast to Coroutine to satisfy mypy, as _receive is awaitable
-                receive_coro: Coroutine[Any, Any, Any] = http_request._receive()  # type: ignore
-                receive_task: Task[Any] = asyncio.create_task(receive_coro)
+                # _receive is a private Starlette/FastAPI method that returns a coroutine
+                receive_obj = http_request  # type: ignore[misc]
+                receive_coro: Coroutine[Any, Any, Dict[str, Any]] = (
+                    receive_obj._receive()
+                )  # type: ignore[misc]
+                receive_task: Task[Dict[str, Any]] = asyncio.create_task(receive_coro)
                 done, pending = await asyncio.wait([receive_task], timeout=0.01)
 
                 if done:
@@ -34,6 +45,8 @@ async def check_client_connection(req_id: str, http_request: Request) -> bool:
                         await receive_task
                     except asyncio.CancelledError:
                         pass
+            except asyncio.CancelledError:
+                raise
             except Exception:
                 # If checking fails, assume disconnected to be safe, or log and continue?
                 # Usually if _receive fails it might mean connection issues.
@@ -44,6 +57,8 @@ async def check_client_connection(req_id: str, http_request: Request) -> bool:
             return False
 
         return True
+    except asyncio.CancelledError:
+        raise
     except Exception:
         return False
 

api_utils/context_init.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 from logging_utils import set_request_id
 from models import ChatCompletionRequest
 
@@ -13,20 +15,23 @@ async def initialize_request_context(
     state.logger.info("开始处理请求...")
     state.logger.info(f"  请求参数 - Model: {request.model}, Stream: {request.stream}")
 
-    context: RequestContext = {
-        "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,
-        "model_id_to_use": None,
-        "needs_model_switching": False,
-    }
+    context: RequestContext = cast(
+        RequestContext,
+        {
+            "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,
+            "model_id_to_use": None,
+            "needs_model_switching": False,
+        },
+    )
 
     return context

api_utils/context_types.py

@@ -1,19 +1,53 @@
-from typing import Any, List, Optional, TypedDict
+import logging
+from asyncio import Future, Lock
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict
 
 from playwright.async_api import Page as AsyncPage
 
+if TYPE_CHECKING:
+    from fastapi import Request
+    from fastapi.responses import JSONResponse
 
-class RequestContext(TypedDict, total=False):
-    logger: Any
-    page: Optional[AsyncPage]
+    from models.chat import ChatCompletionRequest
+
+
+class QueueItem(TypedDict):
+    """Type definition for items in the request queue.
+
+    This defines the structure of each item put into the request_queue,
+    ensuring type safety for queue operations.
+    """
+
+    req_id: str
+    request_data: "ChatCompletionRequest"
+    http_request: "Request"
+    result_future: "Future[JSONResponse]"
+    enqueue_time: float
+    cancelled: bool
+
+
+class RequestContext(TypedDict):
+    """Request context with all keys always present after initialization.
+
+    All keys are required (always exist in the dict) after context_init.py initialization.
+    Optional[] types indicate that the VALUE can be None, not that the key might not exist.
+    """
+
+    # Core components (always set by context_init.py)
+    logger: logging.Logger
+    page: Optional[AsyncPage]  # Value can be None if browser not ready
     is_page_ready: bool
-    parsed_model_list: List[dict]
-    current_ai_studio_model_id: Optional[str]
-    model_switching_lock: Any
-    page_params_cache: dict
-    params_cache_lock: Any
+    parsed_model_list: List[Dict[str, Any]]
+    current_ai_studio_model_id: Optional[str]  # Value can be None initially
+
+    # Locks (always set by server_state)
+    model_switching_lock: Lock
+    page_params_cache: Dict[str, Any]
+    params_cache_lock: Lock
+
+    # Request-specific state (always initialized)
     is_streaming: bool
     model_actually_switched: bool
-    requested_model: Optional[str]
-    model_id_to_use: Optional[str]
+    requested_model: Optional[str]  # Value can be None if not specified
+    model_id_to_use: Optional[str]  # Value set during model analysis
     needs_model_switching: bool

api_utils/dependencies.py

@@ -6,6 +6,8 @@
 from asyncio import Event, Lock, Queue
 from typing import Any, Dict, List, Set
 
+from api_utils.context_types import QueueItem
+
 
 def get_logger() -> logging.Logger:
     from server import logger
@@ -19,7 +21,7 @@ def get_log_ws_manager():
     return log_ws_manager
 
 
-def get_request_queue() -> Queue:
+def get_request_queue() -> "Queue[QueueItem]":
     from server import request_queue
 
     return request_queue

api_utils/mcp_adapter.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 from typing import Any, Dict
@@ -30,6 +31,8 @@ async def execute_mcp_tool(name: str, params: Dict[str, Any]) -> str:
         resp.raise_for_status()
         try:
             data = resp.json()
+        except asyncio.CancelledError:
+            raise
         except Exception:
             data = {"raw": resp.text}
     return json.dumps(data, ensure_ascii=False)
@@ -47,6 +50,8 @@ async def execute_mcp_tool_with_endpoint(
         resp.raise_for_status()
         try:
             data = resp.json()
+        except asyncio.CancelledError:
+            raise
         except Exception:
             data = {"raw": resp.text}
     return json.dumps(data, ensure_ascii=False)

api_utils/model_switching.py

@@ -1,3 +1,5 @@
+import logging
+
 from playwright.async_api import Page as AsyncPage
 
 from api_utils.server_state import state
@@ -19,7 +21,9 @@ async def analyze_model_requirements(
         logger.info(f"请求使用模型: {requested_model_id}")
 
         if parsed_model_list:
-            valid_model_ids = [m.get("id") for m in parsed_model_list]
+            valid_model_ids = [
+                str(m.get("id")) for m in parsed_model_list if m.get("id")
+            ]
             if requested_model_id not in valid_model_ids:
                 from .error_utils import bad_request
 
@@ -50,6 +54,10 @@ async def handle_model_switching(
     model_switching_lock = context["model_switching_lock"]
     model_id_to_use = context["model_id_to_use"]
 
+    # Assert non-None values required for model switching
+    assert page is not None, "Page must be ready for model switching"
+    assert model_id_to_use is not None, "Target model ID must be set"
+
     async with model_switching_lock:
         if state.current_ai_studio_model_id != model_id_to_use:
             logger.info(
@@ -64,19 +72,25 @@ async def handle_model_switching(
                 context["current_ai_studio_model_id"] = model_id_to_use
                 logger.info(f"模型切换成功: {state.current_ai_studio_model_id}")
             else:
+                # Current model ID should exist when switching fails
+                current_model = state.current_ai_studio_model_id or "unknown"
                 await _handle_model_switch_failure(
                     req_id,
                     page,
                     model_id_to_use,
-                    state.current_ai_studio_model_id,
+                    current_model,
                     logger,
                 )
 
     return context
 
 
 async def _handle_model_switch_failure(
-    req_id: str, page: AsyncPage, model_id_to_use: str, model_before_switch: str, logger
+    req_id: str,
+    page: AsyncPage,
+    model_id_to_use: str,
+    model_before_switch: str,
+    logger: logging.Logger,
 ) -> None:
     set_request_id(req_id)
     logger.warning(f"模型切换至 {model_id_to_use} 失败。")

api_utils/page_response.py

@@ -1,4 +1,6 @@
 import asyncio
+import logging
+from typing import Callable
 
 from playwright.async_api import Error as PlaywrightAsyncError
 from playwright.async_api import Page as AsyncPage
@@ -9,7 +11,10 @@
 
 
 async def locate_response_elements(
-    page: AsyncPage, req_id: str, logger, check_client_disconnected
+    page: AsyncPage,
+    req_id: str,
+    logger: logging.Logger,
+    check_client_disconnected: Callable[[str], bool],
 ) -> None:
     """定位响应容器与文本元素，包含超时与错误处理。"""
     set_request_id(req_id)
@@ -26,6 +31,8 @@ async def locate_response_elements(
         from .error_utils import upstream_error
 
         raise upstream_error(req_id, f"定位AI Studio响应元素失败: {locate_err}")
+    except asyncio.CancelledError:
+        raise
     except Exception as locate_exc:
         from .error_utils import server_error