From f50351ee49c12d71ed126d3ab46663e2d4dbb445 Mon Sep 17 00:00:00 2001
From: "Maarten A. Breddels" <maartenbreddels@gmail.com>
Date: Tue, 29 Aug 2023 10:51:34 +0200
Subject: [PATCH] fix: remove websocket object from session on disconnect

---
 solara/server/server.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/solara/server/server.py b/solara/server/server.py
index 2233caf01..ccf834411 100644
--- a/solara/server/server.py
+++ b/solara/server/server.py
@@ -135,6 +135,10 @@ async def app_loop(ws: websocket.WebsocketWrapper, session_id: str, connection_i
             try:
                 message = await ws.receive()
             except websocket.WebSocketDisconnect:
+                try:
+                    context.kernel.session.websockets.remove(ws)
+                except KeyError:
+                    pass
                 logger.debug("Disconnected")
                 return
             t0 = time.time()