diff --git a/app.py b/app.py
index 3e3e5b5..d1cbaac 100644
--- 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()