fix: check python id before deleting UI element

marimo/_plugins/ui/_core/registry.py

@@ -151,17 +151,17 @@ def delete(self, object_id: UIElementId, python_id: int) -> None:
             return
 
         ui_element = self._objects[object_id]()
-        del self._objects[object_id]
-
+        # We guard against UIElement's destructor racing against
+        # registration of another element when a cell re-runs by checking
+        # the Python object id. This isn't perfect because python ids can
+        # be reused ...
         registered_python_id = (
             id(ui_element) if ui_element is not None else None
         )
         if (
             registered_python_id is not None
             and registered_python_id != python_id
         ):
-            # guards against UIElement's destructor racing against
-            # registration of another element when a cell re-runs
             return
 
         try:
@@ -175,3 +175,4 @@ def delete(self, object_id: UIElementId, python_id: int) -> None:
             del self._bindings[object_id]
         if object_id in self._constructing_cells:
             del self._constructing_cells[object_id]
+        del self._objects[object_id]

tests/_plugins/ui/_core/test_registry.py

@@ -168,3 +168,25 @@ async def test_parent_bound_to_view(
     array = k.globals["array"]
     registry = get_context().ui_element_registry
     assert registry.bound_names(array._id) == set(["array", "child"])
+
+
+async def test_dont_delete_element_with_wrong_python_id(
+    k: Kernel, exec_req: ExecReqProvider
+) -> None:
+    await k.run(
+        [
+            exec_req.get("import marimo as mo"),
+            exec_req.get(
+                """
+                s = mo.ui.slider(1, 10)
+                """
+            ),
+        ]
+    )
+    # Make sure that the slider is registered
+    s = k.globals["s"]
+    assert get_context().ui_element_registry.get_object(s._id) == s
+
+    # If the Python id doesn't match, don't delete the object.
+    get_context().ui_element_registry.delete(s._id, -1)
+    assert get_context().ui_element_registry.get_object(s._id) == s