diff --git a/lua/docscribe/core/generator.lua b/lua/docscribe/core/generator.lua
index 9239503..ebbaf9d 100644
--- a/lua/docscribe/core/generator.lua
+++ b/lua/docscribe/core/generator.lua
@@ -73,17 +73,14 @@ function M.generate_docs(function_text, insertion_row, indentation_level)
 
     notification_utils.start_spinner_notification()
     llm_utils.generate_docs(function_text, function(docs, err)
-        is_generating = false
+        vim.schedule(function()
+            is_generating = false
 
-        if not docs then
-            vim.schedule(function()
+            if not docs then
                 notification_utils.replace_spinner_notification("Could not generate docs: " .. err, true)
                 highlight_utils.clear_highlight()
-            end)
-            return
-        end
-
-        vim.schedule(function()
+                return
+            end
             handle_successful_doc_generation(docs, insertion_row, indentation_level)
         end)
     end)
diff --git a/lua/docscribe/ui/notifications.lua b/lua/docscribe/ui/notifications.lua
index 7d5c75c..0c5bb4d 100644
--- a/lua/docscribe/ui/notifications.lua
+++ b/lua/docscribe/ui/notifications.lua
@@ -15,6 +15,10 @@ local spinner_notification_id
 --- Updates the spinner animation by cycling through the `spinner_chars`.
 --- This function updates the spinner message in the notification.
 local function update_spinner()
+    if not spinner_timer then
+        return
+    end
+    
     current_spinner_idx = (current_spinner_idx % #spinner_chars) + 1
 
     spinner_notification_id =
@@ -62,14 +66,12 @@ function M.replace_spinner_notification(message, caught_an_error)
 
     local log_level = caught_an_error and vim.log.levels.ERROR or vim.log.levels.INFO
 
-    if spinner_notification_id then
-        M.docscribe_notify(message, log_level, {
-            timeout = config.get_config("ui").highlight.timeout,
-            hide_from_history = false,
-            replace = spinner_notification_id,
-        })
-        spinner_notification_id = nil
-    end
+    M.docscribe_notify(message, log_level, {
+        timeout = config.get_config("ui").highlight.timeout,
+        hide_from_history = false,
+        replace = spinner_notification_id,
+    })
+    spinner_notification_id = nil
 end
 
 return M