diff --git a/gaphor/core/modeling/presentation.py b/gaphor/core/modeling/presentation.py
index 3b83673ab..233bffa82 100644
--- a/gaphor/core/modeling/presentation.py
+++ b/gaphor/core/modeling/presentation.py
@@ -105,14 +105,7 @@ class Presentation(Matrices, Element, Generic[S]):
             self.parent.matrix_i2c.remove_handler(self._on_matrix_changed)
 
         if diagram := self._original_diagram:
-            connecting_items = [
-                cinfo.item
-                for cinfo in diagram.connections.get_connections(connected=self)
-            ]
-
             diagram.connections.remove_connections_to_item(self)
-            for item in connecting_items:
-                item.unlink()
             self._original_diagram = None
             super().inner_unlink(UnlinkEvent(self, diagram=diagram))
 
diff --git a/gaphor/diagram/connectors.py b/gaphor/diagram/connectors.py
index 3e84357dd..114dde043 100644
--- a/gaphor/diagram/connectors.py
+++ b/gaphor/diagram/connectors.py
@@ -71,7 +71,15 @@ class BaseConnector:
     ) -> None:
         self.element = element
         self.line = line
-        self.diagram: Diagram = element.diagram
+
+        # Try to find the diagram. Any diagram can be None, when called while an element is unlinking
+        assert (
+            element.diagram is line.diagram
+            or element.diagram is None
+            or line.diagram is None
+        )
+        self.diagram: Diagram = element.diagram or line.diagram
+        assert self.diagram
 
     def get_connected(self, handle: Handle) -> Presentation[Element] | None:
         """Get item connected to a handle."""
diff --git a/gaphor/diagram/general/tests/test_comment.py b/gaphor/diagram/general/tests/test_comment.py
index 651063e32..9ed5360b4 100644
--- a/gaphor/diagram/general/tests/test_comment.py
+++ b/gaphor/diagram/general/tests/test_comment.py
@@ -182,8 +182,8 @@ def test_commentline_element_unlink(create, diagram):
 
     assert clazz not in diagram.ownedPresentation
     assert not clazz.diagram
-    assert line not in diagram.ownedPresentation
-    assert not line.diagram
+    assert line in diagram.ownedPresentation
+    assert line.diagram
     assert not comment.subject.annotatedElement
     assert len(clazz_subject.comment) == 0
 
diff --git a/gaphor/diagram/tests/test_presentation.py b/gaphor/diagram/tests/test_presentation.py
index 1ceaad6dc..0c89b1986 100644
--- a/gaphor/diagram/tests/test_presentation.py
+++ b/gaphor/diagram/tests/test_presentation.py
@@ -150,5 +150,5 @@ def test_remove_connected_items_on_unlink(create, diagram):
     class_a_item.subject.unlink()
 
     assert class_a_item not in diagram.ownedPresentation
-    assert association_item not in diagram.ownedPresentation
+    assert association_item in diagram.ownedPresentation
     assert class_b_item in diagram.ownedPresentation
diff --git a/tests/test_undo.py b/tests/test_undo.py
index 34e313668..d0f0eccca 100644
--- a/tests/test_undo.py
+++ b/tests/test_undo.py
@@ -63,7 +63,7 @@ def test_class_association_undo_redo(event_manager, element_factory, undo_manage
     assert undo_manager.can_undo()
 
     for _i in range(3):
-        assert 8 == len(diagram.connections.solver.constraints)
+        assert 9 == len(diagram.connections.solver.constraints)
 
         undo_manager.undo_transaction()