Fix: Restore cursor when updating from outside of tree

lua/gitlab/actions/comment.lua

@@ -41,13 +41,7 @@ local confirm_create_comment = function(text, visual_range, unlinked, discussion
     local body = { discussion_id = discussion_id, reply = text, draft = is_draft }
     job.run_job("/mr/reply", "POST", body, function()
       u.notify("Sent reply!", vim.log.levels.INFO)
-      if is_draft then
-        draft_notes.load_draft_notes(function()
-          discussions.rebuild_view(unlinked)
-        end)
-      else
-        discussions.rebuild_view(unlinked)
-      end
+      discussions.rebuild_view(unlinked)
     end)
     return
   end
@@ -69,8 +63,6 @@ local confirm_create_comment = function(text, visual_range, unlinked, discussion
     return
   end
 
-  vim.print("Here: ", unlinked, discussion_id)
-
   local reviewer_data = reviewer.get_reviewer_data()
   if reviewer_data == nil then
     u.notify("Error getting reviewer data", vim.log.levels.ERROR)
@@ -102,7 +94,7 @@ local confirm_create_comment = function(text, visual_range, unlinked, discussion
     job.run_job("/mr/draft_notes/", "POST", body, function()
       u.notify("Draft reply created!", vim.log.levels.INFO)
       draft_notes.load_draft_notes(function()
-        discussions.rebuild_view(false, true)
+        discussions.rebuild_view(unlinked)
       end)
     end)
     return

lua/gitlab/actions/discussions/init.lua

@@ -139,7 +139,7 @@ M.open = function(callback)
   local current_window = vim.api.nvim_get_current_win() -- Save user's current window in case they switched while content was loading
   vim.api.nvim_set_current_win(M.split.winid)
 
-  common.switch_can_edit_bufs(true, M.linked_bufnr, M.unliked_bufnr)
+  common.switch_can_edit_bufs(true, M.linked_bufnr, M.unlinked_bufnr)
   M.rebuild_discussion_tree()
   M.rebuild_unlinked_discussion_tree()
 
@@ -432,6 +432,9 @@ M.rebuild_discussion_tree = function()
   if M.linked_bufnr == nil then
     return
   end
+
+  local current_node = discussions_tree.get_node_at_cursor(M.discussion_tree, M.last_node_at_cursor)
+
   local expanded_node_ids = M.gather_expanded_node_ids(M.discussion_tree)
   common.switch_can_edit_bufs(true, M.linked_bufnr, M.unlinked_bufnr)
 
@@ -447,12 +450,14 @@ M.rebuild_discussion_tree = function()
     bufnr = M.linked_bufnr,
     prepare_node = tree_utils.nui_tree_prepare_node,
   })
+
   -- Re-expand already expanded nodes
   for _, id in ipairs(expanded_node_ids) do
     tree_utils.open_node_by_id(discussion_tree, id)
   end
-
   discussion_tree:render()
+  discussions_tree.restore_cursor_position(M.split.winid, discussion_tree, current_node)
+
   M.set_tree_keymaps(discussion_tree, M.linked_bufnr, false)
   M.discussion_tree = discussion_tree
   common.switch_can_edit_bufs(false, M.linked_bufnr, M.unlinked_bufnr)
@@ -466,6 +471,9 @@ M.rebuild_unlinked_discussion_tree = function()
   if M.unlinked_bufnr == nil then
     return
   end
+
+  local current_node = discussions_tree.get_node_at_cursor(M.unlinked_discussion_tree, M.last_node_at_cursor)
+
   local expanded_node_ids = M.gather_expanded_node_ids(M.unlinked_discussion_tree)
   common.switch_can_edit_bufs(true, M.linked_bufnr, M.unlinked_bufnr)
   vim.api.nvim_buf_set_lines(M.unlinked_bufnr, 0, -1, false, {})
@@ -487,6 +495,7 @@ M.rebuild_unlinked_discussion_tree = function()
     tree_utils.open_node_by_id(unlinked_discussion_tree, id)
   end
   unlinked_discussion_tree:render()
+  discussions_tree.restore_cursor_position(M.split.winid, unlinked_discussion_tree, current_node)
 
   M.set_tree_keymaps(unlinked_discussion_tree, M.unlinked_bufnr, true)
   M.unlinked_discussion_tree = unlinked_discussion_tree
@@ -535,6 +544,14 @@ M.create_split_and_bufs = function()
     buffer = linked_bufnr,
     callback = function()
       M.last_row, M.last_column = unpack(vim.api.nvim_win_get_cursor(0))
+      M.last_node_at_cursor = M.discussion_tree and M.discussion_tree:get_node() or nil
+    end,
+  })
+
+  vim.api.nvim_create_autocmd("WinLeave", {
+    buffer = unlinked_bufnr,
+    callback = function()
+      M.last_node_at_cursor = M.unlinked_discussion_tree and M.unlinked_discussion_tree:get_node() or nil
     end,
   })
 

lua/gitlab/actions/discussions/tree.lua

@@ -422,12 +422,36 @@ M.toggle_nodes = function(winid, tree, unlinked, opts)
   M.restore_cursor_position(winid, tree, current_node, root_node)
 end
 
+-- Get current node for restoring cursor position
+---@param tree NuiTree The inline discussion tree or the unlinked discussion tree
+---@param last_node NuiTree.Node|nil The last active discussion tree node in case we are not in any of the discussion trees
+M.get_node_at_cursor = function(tree, last_node)
+  if tree == nil then
+    return
+  end
+  if vim.api.nvim_get_current_win() == vim.fn.win_findbuf(tree.bufnr)[1] then
+    return tree:get_node()
+  else
+    return last_node
+  end
+end
+
 ---Restore cursor position to the original node if possible
+---@param winid integer Window number of the discussions split
+---@param tree NuiTree The inline discussion tree or the unlinked discussion tree
+---@param original_node NuiTree.Node|nil The last node with the cursor
+---@param root_node NuiTree.Node|nil The root node of the last node with the cursor
 M.restore_cursor_position = function(winid, tree, original_node, root_node)
+  if original_node == nil or tree == nil then
+    return
+  end
   local _, line_number = tree:get_node("-" .. tostring(original_node.id))
-  -- If current_node is has been collapsed, get line number of root node instead
-  if line_number == nil and root_node then
-    _, line_number = tree:get_node("-" .. tostring(root_node.id))
+  -- If current_node has been collapsed, try to get line number of root node instead
+  if line_number == nil then
+    root_node = root_node and root_node or common.get_root_node(tree, original_node)
+    if root_node ~= nil then
+      _, line_number = tree:get_node("-" .. tostring(root_node.id))
+    end
   end
   if line_number ~= nil then
     vim.api.nvim_win_set_cursor(winid, { line_number, 0 })