From e3af38340804dae05a6beff2c49b24b5ae6d2796 Mon Sep 17 00:00:00 2001
From: "nhat.vo" <tonystark0243@gmail.com>
Date: Tue, 21 Feb 2023 13:30:39 +0100
Subject: [PATCH] added cleanup to session load event

---
 lua/neo-tree/events/init.lua | 1 +
 lua/neo-tree/setup/init.lua  | 8 ++++++++
 lua/neo-tree/ui/renderer.lua | 9 +++++----
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/lua/neo-tree/events/init.lua b/lua/neo-tree/events/init.lua
index cbd108ed..2a73b49c 100644
--- a/lua/neo-tree/events/init.lua
+++ b/lua/neo-tree/events/init.lua
@@ -41,6 +41,7 @@ local M = {
   VIM_LEAVE = "vim_leave",
   VIM_RESIZED = "vim_resized",
   VIM_CURSOR_MOVED = "vim_cursor_moved",
+  VIM_AFTER_SESSION_LOAD = "vim_after_session_load",
 }
 
 M.define_autocmd_event = function(event_name, autocmds, debounce_frequency, seed_fn, nested)
diff --git a/lua/neo-tree/setup/init.lua b/lua/neo-tree/setup/init.lua
index ff9cee41..74c9ebb5 100644
--- a/lua/neo-tree/setup/init.lua
+++ b/lua/neo-tree/setup/init.lua
@@ -76,6 +76,7 @@ local define_events = function()
   events.define_autocmd_event(events.VIM_WIN_CLOSED, { "WinClosed" })
   events.define_autocmd_event(events.VIM_COLORSCHEME, { "ColorScheme" }, 0)
   events.define_autocmd_event(events.VIM_CURSOR_MOVED, { "CursorMoved" }, 100)
+  events.define_autocmd_event(events.VIM_AFTER_SESSION_LOAD, { "SessionLoadPost" }, 200)
   events.define_autocmd_event(events.GIT_EVENT, { "User FugitiveChanged" }, 100)
   events.define_event(events.GIT_STATUS_CHANGED, { debounce_frequency = 0 })
   events_setup = true
@@ -93,6 +94,13 @@ local define_events = function()
       require("neo-tree.ui.renderer").update_floating_window_layouts()
     end,
   })
+
+  events.subscribe({
+    event = events.VIM_AFTER_SESSION_LOAD,
+    handler = function()
+      require("neo-tree.ui.renderer").clean_invalid_neotree_buffers(true)
+    end,
+  })
 end
 
 local prior_window_options = {}
diff --git a/lua/neo-tree/ui/renderer.lua b/lua/neo-tree/ui/renderer.lua
index a8ae420e..1c8cdd11 100644
--- a/lua/neo-tree/ui/renderer.lua
+++ b/lua/neo-tree/ui/renderer.lua
@@ -27,10 +27,11 @@ local update_floating_windows = function()
   floating_windows = valid_windows
 end
 
--- clean up neotree buffers created by :mksession
 local cleaned_up = false
-local clean_neotree_buffers = function()
-  if cleaned_up then
+---Clean up invalid neotree buffers (e.g after a session restore)
+---@param force boolean if true, force cleanup. Otherwise only cleanup once
+M.clean_invalid_neotree_buffers = function(force)
+  if cleaned_up and not force then
     return
   end
 
@@ -839,7 +840,7 @@ create_window = function(state)
     tabnr = state.tabnr,
   }
   events.fire_event(events.NEO_TREE_WINDOW_BEFORE_OPEN, event_args)
-  clean_neotree_buffers()
+  M.clean_invalid_neotree_buffers(false)
 
   if state.current_position == "float" then
     state.force_float = nil