Merge pull request #221 from Maxteabag/worktree-filter-cursor-fix

Land filter-accept cursor on the matched node after snapshot restore

Author: Maxteabag

sqlit/domains/explorer/ui/mixins/tree_filter.py

@@ -109,17 +109,61 @@ def action_tree_filter_close(self: TreeFilterMixinHost) -> None:
 
     def action_tree_filter_accept(self: TreeFilterMixinHost) -> None:
         """Accept current filter selection, close filter, and activate the node."""
-        # Store current match before closing
-        current_node = None
-        if self._tree_filter_matches and self._tree_filter_match_index < len(self._tree_filter_matches):
+        # Remember the match's *data* (not the node reference) before closing.
+        # Closing the filter rebuilds the tree from the snapshot taken at
+        # filter-open time, which replaces every node object — so the
+        # reference we captured here would be stale after close. The data
+        # payload, however, is the same object on both old and new nodes
+        # (we pass it through unchanged in _restore_node_under), so we can
+        # re-locate the match by identity.
+        matched_data: Any = None
+        if (
+            self._tree_filter_matches
+            and self._tree_filter_match_index < len(self._tree_filter_matches)
+        ):
             current_node = self._tree_filter_matches[self._tree_filter_match_index]
+            if current_node and current_node.data:
+                matched_data = current_node.data
 
         # Close the filter
         self.action_tree_filter_close()
 
-        # Activate the selected node (connect to server, expand folder, etc.)
-        if current_node and current_node.data:
-            self._activate_tree_node(current_node)
+        if matched_data is None:
+            return
+
+        fresh_node = self._find_node_by_data(matched_data)
+        if fresh_node is None:
+            return
+
+        # Textual's Tree.move_cursor reads `node._line`, which is set during
+        # the next layout pass — not when the node is `add()`-ed. Since the
+        # snapshot restore that just ran in action_tree_filter_close added
+        # all fresh nodes synchronously, calling move_cursor right now sees
+        # stale `_line` values and parks the cursor on the wrong row.
+        # Defer the move (and the activation) until after the next refresh.
+        call_after = getattr(self, "call_after_refresh", None)
+        if callable(call_after):
+            call_after(lambda: self._select_and_activate_after_refresh(fresh_node))
+        else:
+            # Synchronous fallback (used in unit tests with a mock host).
+            self._select_and_activate_after_refresh(fresh_node)
+
+    def _select_and_activate_after_refresh(self: TreeFilterMixinHost, node: Any) -> None:
+        try:
+            self.object_tree.move_cursor(node)
+        except Exception:
+            pass
+        self._activate_tree_node(node)
+
+    def _find_node_by_data(self: TreeFilterMixinHost, data: Any) -> Any | None:
+        """Locate the node in the current tree whose `.data` is `data`."""
+        stack = [self.object_tree.root]
+        while stack:
+            node = stack.pop()
+            if node.data is data:
+                return node
+            stack.extend(node.children)
+        return None
 
     def action_tree_filter_next(self: TreeFilterMixinHost) -> None:
         """Move to next filter match."""

tests/ui/explorer/test_filter_accept_cursor_pilot.py

@@ -0,0 +1,80 @@
+"""Pilot-driven test for cursor position after accepting a filter match.
+
+Runs the real Textual app + Tree widget via `app.run_test()` and presses
+real keys. Asserts that pressing Enter on a filtered match leaves the
+cursor on that match in the rebuilt tree (the failure mode the user
+sees is that it lands somewhere else — e.g. the top node).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from sqlit.domains.shell.app.main import SSMSTUI
+
+from ..mocks import (
+    MockConnectionStore,
+    MockSettingsStore,
+    build_test_services,
+    create_test_connection,
+)
+
+
+def _connection_name_of(node) -> str | None:
+    data = getattr(node, "data", None)
+    if data is None:
+        return None
+    config = getattr(data, "config", None)
+    if config is None:
+        return None
+    return getattr(config, "name", None)
+
+
+@pytest.mark.asyncio
+async def test_filter_accept_lands_cursor_on_matched_connection():
+    """Repro of the issue the user demonstrated: after `/` then typing a
+    substring and pressing Enter, the explorer cursor should be on the
+    matched connection — not on the top node."""
+    connections = [
+        create_test_connection("alpha", "sqlite"),
+        create_test_connection("test-server", "sqlite"),
+        create_test_connection("production", "sqlite"),
+    ]
+    services = build_test_services(
+        connection_store=MockConnectionStore(connections),
+        settings_store=MockSettingsStore({"theme": "tokyo-night"}),
+    )
+    app = SSMSTUI(services=services)
+
+    async with app.run_test(size=(120, 40)) as pilot:
+        await pilot.pause()
+
+        # Focus the explorer so key events reach the tree filter mixin.
+        app.action_focus_explorer()
+        await pilot.pause()
+        assert app.object_tree.has_focus
+
+        # Open the filter, type a substring of one connection, accept.
+        await pilot.press("slash")
+        await pilot.pause()
+        for ch in "test":
+            await pilot.press(ch)
+            await pilot.pause()
+
+        # Sanity: a single match was found.
+        assert len(app._tree_filter_matches) == 1, (
+            f"expected exactly one match for 'test'; "
+            f"got {len(app._tree_filter_matches)}"
+        )
+
+        await pilot.press("enter")
+        await pilot.pause()
+
+        cursor = app.object_tree.cursor_node
+        assert cursor is not None, "cursor unexpectedly None after accept"
+        name = _connection_name_of(cursor)
+        assert name == "test-server", (
+            "Cursor landed on the wrong element after accepting the filter "
+            f"match. Expected 'test-server', got node with connection "
+            f"name={name!r} (cursor.label={getattr(cursor, 'label', None)!r})."
+        )

tests/ui/explorer/test_tree_filter_search.py

@@ -64,9 +64,29 @@ def __init__(self):
     def select_node(self, node: MockTreeNode) -> None:
         self.selected_node = node
 
+    def move_cursor(self, node: MockTreeNode) -> None:
+        self.selected_node = node
+
     def focus(self) -> None:
         self.has_focus = True
 
+    def is_node_in_tree(self, node: MockTreeNode | None) -> bool:
+        """Walk the tree to check if a node reference is still attached.
+
+        Mirrors Textual's actual behavior: a cursor reference pointing at a
+        node that was `child.remove()`d is stale — the widget no longer
+        renders that node.
+        """
+        if node is None:
+            return False
+        stack = [self.root]
+        while stack:
+            current = stack.pop()
+            if current is node:
+                return True
+            stack.extend(current.children)
+        return False
+
 
 class MockFilterInput:
     """Mock TreeFilterInput capturing the last set_filter call."""
@@ -384,3 +404,52 @@ def test_filter_finds_lazy_loaded_tables_in_multi_db_mode(self):
         assert "cs_ticket" in matched, (
             f"Issue #141: filter must find lazy-loaded tables; got {matched}"
         )
+
+
+class TestCursorPositionAfterFilterAccept:
+    """Pressing Enter on a filter match should leave the cursor on that
+    same match in the rebuilt tree — not on a stale node reference (the
+    pre-snapshot match object has been removed and replaced by a fresh
+    node when the tree was restored)."""
+
+    def _open_filter(self, host: _FilterHost) -> None:
+        TreeFilterMixin.action_tree_filter(host)  # type: ignore[arg-type]
+
+    def _type(self, host: _FilterHost, text: str) -> None:
+        for ch in text:
+            host._tree_filter_text += ch
+            TreeFilterMixin._update_tree_filter(host)  # type: ignore[arg-type]
+
+    def test_cursor_stays_on_matched_node_after_accept(self):
+        host = _FilterHost(["alpha", "test-server", "production"])
+
+        self._open_filter(host)
+        self._type(host, "test")
+
+        # We have exactly one match: 'test-server'
+        assert len(host._tree_filter_matches) == 1
+        matched_label = host._tree_filter_matches[0].data.get_label_text()
+        assert matched_label == "test-server"
+
+        # Accept (this closes the filter and restores the tree from snapshot)
+        TreeFilterMixin.action_tree_filter_accept(host)  # type: ignore[arg-type]
+
+        cursor = host.object_tree.selected_node
+
+        # 1. The cursor must point at a node that is still in the tree —
+        #    not at a removed/stale Python object from before the restore.
+        assert host.object_tree.is_node_in_tree(cursor), (
+            "After accept, cursor references a node that's no longer in "
+            "the tree (stale reference left over from the filter session). "
+            f"cursor: {cursor!r}"
+        )
+
+        # 2. That node should correspond to the match the user accepted.
+        assert cursor is not None
+        assert (
+            cursor.data is not None
+            and cursor.data.get_label_text() == "test-server"
+        ), (
+            "Cursor ended up on the wrong node after accept. "
+            f"Expected 'test-server', got: {cursor.data.get_label_text() if cursor.data else None}"
+        )