Litestar框架中WebSocket连接生命周期的正确实现方式

在Litestar框架中，WebSocket监听器的connection_lifespan参数是一个强大的功能，它允许开发者控制WebSocket连接的整个生命周期。然而，官方文档中的示例存在一个关键的技术细节需要特别注意。

生命周期管理器的本质

WebSocket连接的生命周期管理器实际上是一个异步上下文管理器，这意味着它必须满足特定的结构要求：

1. 必须使用@contextlib.asynccontextmanager装饰器
2. 必须是一个异步生成器函数
3. 必须使用yield语句来划分连接的生命周期阶段

正确实现模式

以下是符合Litestar要求的WebSocket连接生命周期管理器的标准实现：

from contextlib import asynccontextmanager
from typing import AsyncGenerator

from litestar import WebSocket

@asynccontextmanager
async def proper_lifespan(socket: WebSocket) -> AsyncGenerator[None, None]:
    # 连接建立时的初始化代码
    print(f"Connection established from {socket.client}")
    
    try:
        yield  # 这里将控制权交还给连接处理逻辑
    finally:
        # 连接关闭时的清理代码
        print(f"Connection closed for {socket.client}")

类型提示的重要性

为了通过mypy类型检查，生命周期管理器的类型注解必须准确：

• 返回类型必须是AsyncGenerator[None, None]
• 使用@asynccontextmanager装饰器确保类型兼容性

常见错误模式

开发者容易犯的几个典型错误：

1. 忘记使用@asynccontextmanager装饰器
2. 将生命周期管理器写成普通异步函数而非生成器
3. 缺少yield语句
4. 类型注解不完整或不正确

完整示例

结合Litestar应用的完整实现：

from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator

from litestar import Litestar, WebSocket, websocket_listener

@asynccontextmanager
async def connection_lifespan(socket: WebSocket) -> AsyncGenerator[None, None]:
    print("Connection opened")
    try:
        yield
    finally:
        print("Connection closed")

@websocket_listener("/", connection_lifespan=connection_lifespan)
async def websocket_handler(socket: WebSocket, data: Any) -> None:
    await socket.send_json({"message": "Hello from server"})

app = Litestar([websocket_handler])