From 95c6ccbe83718ae908ad3cee0e9b0f2edc3b2def Mon Sep 17 00:00:00 2001
From: xodn348 <xodn348@tamu.edu>
Date: Thu, 21 May 2026 09:23:21 +0000
Subject: [PATCH] fix(registry): pass element id, not weakref id, to delete()
 on re-register
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

In UIElementRegistry.register(), when an object_id is already known, the
old element is deleted via:

    self.delete(object_id, id(self._objects[object_id]))

self._objects[object_id] is a weakref.ref wrapper, so id() here yields the
wrapper's address — not the element's. In delete(), the guard compares that
python_id against id(element) (dereferenced). They never match, so delete()
always returns early without calling ctx.function_registry.delete().

Fix: dereference the weakref first:

    self.delete(object_id, id(self._objects[object_id]()))

If the referent has already been GC'd (returns None), id(None) is passed;
delete()'s guard sets registered_python_id=None and the element is still
removed correctly.

Fixes #9551
---
 marimo/_plugins/ui/_core/registry.py     |  2 +-
 tests/_plugins/ui/_core/test_registry.py | 48 ++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/marimo/_plugins/ui/_core/registry.py b/marimo/_plugins/ui/_core/registry.py
index 38803488230..0e133a94075 100644
--- a/marimo/_plugins/ui/_core/registry.py
+++ b/marimo/_plugins/ui/_core/registry.py
@@ -37,7 +37,7 @@ def register(
             # on cell re-run, a UI element may be (re)-registered before
             # its destructor was called, so manually delete the old element
             # here
-            self.delete(object_id, id(self._objects[object_id]))
+            self.delete(object_id, id(self._objects[object_id]()))
         self._objects[object_id] = weakref.ref(ui_element)
         assert execution_context is not None
         self._constructing_cells[object_id] = execution_context.cell_id
diff --git a/tests/_plugins/ui/_core/test_registry.py b/tests/_plugins/ui/_core/test_registry.py
index e1491927aa3..283ed478da0 100644
--- a/tests/_plugins/ui/_core/test_registry.py
+++ b/tests/_plugins/ui/_core/test_registry.py
@@ -1,9 +1,14 @@
 # Copyright 2026 Marimo. All rights reserved.
 from __future__ import annotations
 
+import weakref
+from unittest.mock import MagicMock, patch
+
 from marimo import ui
+from marimo._plugins.ui._core.registry import UIElementRegistry
 from marimo._runtime.context import get_context
 from marimo._runtime.runtime import Kernel
+from marimo._types.ids import UIElementId
 from tests.conftest import ExecReqProvider
 
 
@@ -170,6 +175,49 @@ async def test_parent_bound_to_view(
     assert registry.bound_names(array._id) == {"array", "child"}
 
 
+def test_register_passes_element_id_not_weakref_id() -> None:
+    """register() must call delete() with id(element), not id(weakref.ref).
+
+    Regression test for the bug where register() passed id(self._objects[oid])
+    — the weakref wrapper — rather than id(self._objects[oid]()) — the element.
+    Because delete() derives registered_python_id from id(element), passing the
+    weakref id always fails the guard check and causes delete() to silently
+    return without cleaning up function-registry entries for the old element.
+    """
+
+    class _FakeElement:
+        _lens = None
+
+    elem1 = _FakeElement()
+    elem2 = _FakeElement()
+    oid = UIElementId("ui-test-weakref-id")
+
+    registry = UIElementRegistry()
+    registry._objects[oid] = weakref.ref(elem1)
+
+    delete_calls: list[int] = []
+
+    def _tracking_delete(_object_id: UIElementId, python_id: int) -> None:
+        delete_calls.append(python_id)
+
+    registry.delete = _tracking_delete  # type: ignore[method-assign]
+
+    mock_ctx = MagicMock()
+    mock_ctx.execution_context = MagicMock()
+
+    with patch(
+        "marimo._plugins.ui._core.registry.get_context",
+        return_value=mock_ctx,
+    ):
+        registry.register(oid, elem2)  # type: ignore[arg-type]
+
+    assert len(delete_calls) == 1
+    assert delete_calls[0] == id(elem1), (
+        f"delete() must receive id(element)={id(elem1)}, "
+        f"not id(weakref)={id(registry._objects.get(oid, 'gone'))}"
+    )
+
+
 async def test_dont_delete_element_with_wrong_python_id(
     k: Kernel, exec_req: ExecReqProvider
 ) -> None: