handle websocket disconnect

2026-01-22 10:22:36 +01:00
parent 48d8ce9276
commit 867b029e1c
1 changed files with 10 additions and 5 deletions
--- a/app.py
+++ b/app.py
@@ -12,6 +12,7 @@ from fastapi import (
    FastAPI,
    Request,
    WebSocket,
+    WebSocketDisconnect,
    Depends,
    HTTPException,
    status,
@@ -95,10 +96,14 @@ async def websocket_endpoint(websocket: WebSocket):
        websocket: The WebSocket connection.
    """
    await websocket.accept()
-    while True:
-        data = await websocket.receive_text()
+    try:
+        while True:
+            data = await websocket.receive_text()

-        async for chunk in generate(data):
-            await websocket.send_text(chunk)
+            async for chunk in generate(data):
+                await websocket.send_text(chunk)

-        await websocket.send_text("<<END>>")
+            await websocket.send_text("<<END>>")
+    except WebSocketDisconnect:
+        logger.info("Client disconnected")
+        await websocket.close()