diff --git a/README.md b/README.md
index 8771d5c..a3dd0d1 100644
--- a/README.md
+++ b/README.md
@@ -89,11 +89,11 @@ require("persisted").setup({
   autoload = false, -- automatically load the session for the cwd on Neovim startup
   allowed_dirs = nil, -- table of dirs that the plugin will auto-save and auto-load from
   ignored_dirs = nil, -- table of dirs that are ignored when auto-saving and auto-loading
-  before_save = function() end, -- function to run before the session is saved to disk
-  after_save = function() end, -- function to run after the session is saved to disk
+  before_save = nil, -- function to run before the session is saved to disk
+  after_save = nil, -- function to run after the session is saved to disk
   telescope = { -- options for the telescope extension
-    before_source = function(session) end, -- function to run before the session is sourced via telescope
-    after_source = function(session) end, -- function to run after the session is sourced via telescope
+    before_source = nil, -- function to run before the session is sourced via telescope
+    after_source = nil, -- function to run after the session is sourced via telescope
   },
 })
 ```
@@ -243,7 +243,7 @@ The plugin allows for *before* and *after* callbacks to be used when sourcing a
 ```lua
 require("persisted").setup({
   telescope = {
-    before_source = function()
+    before_source = function(session)
       -- Close all open buffers
       pcall(vim.cmd, "bufdo bwipeout")
     end,
@@ -254,10 +254,8 @@ require("persisted").setup({
   },
 })
 ```
-
-A *session* table is exposed to the callback functions and has the following properties:
+The callbacks must accept a *session* parameter which is a table that has the following properties:
 * name - The filename of the session.
-* branch - The git branch of the session; *and*
 * file_path - The file path to the session.
 
 ## :rocket: Usage
diff --git a/lua/persisted/config.lua b/lua/persisted/config.lua
index dc16b14..f0e736a 100644
--- a/lua/persisted/config.lua
+++ b/lua/persisted/config.lua
@@ -9,12 +9,12 @@ local defaults = {
   autoload = false, -- automatically load the session for the cwd on Neovim startup
   allowed_dirs = nil, -- table of dirs that the plugin will auto-save and auto-load from
   ignored_dirs = nil, -- table of dirs that are ignored for auto-saving and auto-loading
-  before_save = function() end, -- function to run before the session is saved to disk
-  after_save = function() end, -- function to run after the session is saved to disk
-  after_source = function() end, -- function to run after the session is sourced
+  before_save = nil, -- function to run before the session is saved to disk
+  after_save = nil, -- function to run after the session is saved to disk
+  after_source = nil, -- function to run after the session is sourced
   telescope = { -- options for the telescope extension
-    before_source = function(session) end, -- function to run before the session is sourced via telescope
-    after_source = function(session) end, -- function to run after the session is sourced via telescope
+    before_source = nil, -- function to run before the session is sourced via telescope
+    after_source = nil, -- function to run after the session is sourced via telescope
   },
 }
 
diff --git a/lua/persisted/init.lua b/lua/persisted/init.lua
index 151de0e..fea50ce 100644
--- a/lua/persisted/init.lua
+++ b/lua/persisted/init.lua
@@ -102,7 +102,7 @@ function M.load(opt)
   local session = opt.last and get_last() or get_current()
 
   if session and vim.fn.filereadable(session) ~= 0 then
-    utils.load_session(session, _, config.options.after_source())
+    utils.load_session(session, config.options.before_save, config.options.after_save)
   end
 
   if config.options.autosave and (allow_dir() and not ignore_dir()) then
@@ -140,12 +140,16 @@ end
 ---Save the session to disk
 ---@return nil
 function M.save()
-  config.options.before_save()
+  if type(config.options.before_save) == "function" then
+    config.options.before_save()
+  end
 
   vim.cmd("mks! " .. e(get_current()))
   vim.g.persisting = true
 
-  config.options.after_save()
+  if type(config.options.after_save) == "function" then
+    config.options.after_save()
+  end
 end
 
 ---Delete the current session from disk
diff --git a/lua/persisted/utils.lua b/lua/persisted/utils.lua
index f130686..8f46f20 100644
--- a/lua/persisted/utils.lua
+++ b/lua/persisted/utils.lua
@@ -1,4 +1,5 @@
 local M = {}
+local e = vim.fn.fnameescape
 
 --- Split a string into a table
 ---@param input string
@@ -68,7 +69,7 @@ function M.load_session(session, before_callback, after_callback)
     if type(before_callback) == "function" then
       before_callback()
     end
-    local ok, result = pcall(vim.cmd, "source " .. vim.fn.fnameescape(session))
+    local ok, result = pcall(vim.cmd, "source " .. e(session))
     if not ok then
       return M.echoerr("Error loading the session! ", result)
     end
diff --git a/lua/telescope/_extensions/persisted.lua b/lua/telescope/_extensions/persisted.lua
index a821463..3b2749e 100644
--- a/lua/telescope/_extensions/persisted.lua
+++ b/lua/telescope/_extensions/persisted.lua
@@ -28,11 +28,11 @@ local function search_sessions(opts)
         local session = action_state.get_selected_entry()
         actions.close(prompt_bufnr)
 
-        utils.load_session(session.file_path, function()
-          return config.telescope.before_source(session)
-        end, function()
-          return config.telescope.after_source(session)
-        end)
+        utils.load_session(
+          session.file_path,
+          config.telescope.before_source and config.telescope.before_source(session) or _,
+          config.telescope.after_source and config.telescope.after_source(session) or _
+        )
       end)
       return true
     end,