如何在Django/Django-Ninja中检测StreamingHttpResponse（SSE）的客户端断开？

问题描述 #

我们希望在应用中添加Server Side Events（SSE）功能，但当前的while True循环会无限运行，即便客户端断开连接，任务仍不会停止。示例代码运行后会持续打印“Working...”直到服务器关闭，完全不受客户端连接状态影响。

可复现代码：

@api.get("/events")
async def get_events(request: ASGIRequest) -> StreamingHttpResponse:
    async def generate_message() -> AsyncGenerator[str, None]:
        while True:
            await asyncio.sleep(1)
            yield 'data: {"test": "testvalue"}\n\n'
            print("Working...")

    return StreamingHttpResponse(
        generate_message(),
        content_type="text/event-stream",
        headers={
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )

我们已尝试返回简单HttpResponse、将代码放入try..except块、为StreamingHttpResponse添加close方法，但均无效。

环境版本：

• django-ninja==0.22.2
• django==4.2.13
• uvicorn[standard]==0.29.0

解决方案 #

核心思路是捕获客户端断开时触发的异常，以此终止循环。在ASGI环境中，客户端断开连接后，服务器尝试向客户端发送数据（即yield操作）会触发asyncio.CancelledError、BrokenPipeError或ConnectionResetError，只需在生成器中捕获这些异常即可停止循环。

修改后的代码：

import asyncio

@api.get("/events")
async def get_events(request: ASGIRequest) -> StreamingHttpResponse:
    async def generate_message() -> AsyncGenerator[str, None]:
        try:
            while True:
                await asyncio.sleep(1)
                yield 'data: {"test": "testvalue"}\n\n'
                print("Working...")
        except (asyncio.CancelledError, BrokenPipeError, ConnectionResetError):
            print("客户端已断开连接，停止生成事件")
            return

    return StreamingHttpResponse(
        generate_message(),
        content_type="text/event-stream",
        headers={
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        },
    )

原理说明 #

• 当客户端主动断开连接，Uvicorn等ASGI服务器会取消对应的异步任务，触发asyncio.CancelledError
• 若客户端连接异常中断，可能会触发BrokenPipeError或ConnectionResetError
• 将while True循环放入try块，捕获上述异常后直接返回，即可终止生成器，停止后续的任务执行

内容的提问来源于stack exchange，提问作者AlTosterino