fix: make page_close on close beacon thread safe

Since we check for, and get the kernel on two different lines, we
can trigger a KeyError when the close beacon is handled.

Calling .get on the dict will make this thread safe.

Author: maartenbreddels

solara/server/flask.py

@@ -162,8 +162,8 @@ def kernels_connection(ws: simple_websocket.Server, kernel_id: str, name: str):
 @blueprint.route("/_solara/api/close/<kernel_id>", methods=["GET", "POST"])
 def close(kernel_id: str):
     page_id = request.args["session_id"]
-    if kernel_id in kernel_context.contexts:
-        context = kernel_context.contexts[kernel_id]
+    context = kernel_context.contexts.get(kernel_id, None)
+    if context is not None:
         context.page_close(page_id)
     return ""
 

solara/server/starlette.py

@@ -311,8 +311,8 @@ async def run_wrapper():
 def close(request: Request):
     kernel_id = request.path_params["kernel_id"]
     page_id = request.query_params["session_id"]
-    if kernel_id in kernel_context.contexts:
-        context = kernel_context.contexts[kernel_id]
+    context = kernel_context.contexts.get(kernel_id, None)
+    if context is not None:
         context.page_close(page_id)
     response = HTMLResponse(content="", status_code=200)
     return response