fix: safer check before removing session (#3838)

With multiple clients. It's possible they try to remove each other at
the same time from different callsites, so this adds an extra check

marimo/_ast/app.py

@@ -4,6 +4,7 @@
 import ast
 import base64
 import inspect
+import sys
 import threading
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
@@ -25,7 +26,10 @@
 )
 from uuid import uuid4
 
-from typing_extensions import ParamSpec, TypeAlias
+if sys.version_info < (3, 10):
+    from typing_extensions import ParamSpec, TypeAlias
+else:
+    from typing import ParamSpec, TypeAlias
 
 from marimo import _loggers
 from marimo._ast.cell import Cell, CellConfig, CellId_t, CellImpl

marimo/_ast/cell_manager.py

@@ -5,6 +5,7 @@
 import os
 import random
 import string
+import sys
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -13,7 +14,10 @@
     TypeVar,
 )
 
-from typing_extensions import ParamSpec, TypeAlias
+if sys.version_info < (3, 10):
+    from typing_extensions import ParamSpec, TypeAlias
+else:
+    from typing import ParamSpec, TypeAlias
 
 from marimo._ast.cell import Cell, CellConfig, CellId_t
 from marimo._ast.compiler import cell_factory, toplevel_cell_factory

marimo/_server/sessions.py

@@ -1035,7 +1035,7 @@ def maybe_resume_session(
             # Set new session and remove old session
             self.sessions[new_session_id] = session
             # If the ID is the same, we don't need to delete the old session
-            if new_session_id != session_id:
+            if new_session_id != session_id and session_id in self.sessions:
                 del self.sessions[session_id]
             return session
 
@@ -1088,7 +1088,8 @@ def close_session(self, session_id: SessionId) -> bool:
             )
 
         session.close()
-        del self.sessions[session_id]
+        if session_id in self.sessions:
+            del self.sessions[session_id]
         return True
 
     def close_all_sessions(self) -> None:

tests/snapshots/api.txt

@@ -32,6 +32,7 @@ hstack
 icon
 iframe
 image
+import_guard
 islands
   MarimoIslandGenerator
   MarimoIslandStub