clean up UI on drive viewer

maildir_gtd/app.py

@@ -513,6 +513,7 @@ class EmailViewerApp(App):
             self.fetch_envelopes()
 
     async def action_archive(self) -> None:
+        # Remove from all data structures
         self.all_envelopes = [item for item in self.all_envelopes if item and item.get("id") != self.current_message_id]
         self.envelope_map.pop(self.current_message_id, None)
         self.envelope_index_map = {index: id for index, id in self.envelope_index_map.items() if id != self.current_message_id}
@@ -523,13 +524,30 @@ class EmailViewerApp(App):
             k: v for k, v in self.message_body_cache.items() if k != self.current_message_id
         }
         self.total_messages = len(self.message_metadata)
+
+        # Perform archive operation
         worker = archive_current(self)
         await worker.wait()
-        newmsg = self.all_envelopes[self.current_message_index]
-        if newmsg.get("type") == "header":
-            newmsg = self.all_envelopes[self.current_message_index + 1]
-            return
-        self.show_message(newmsg["id"])
+
+        # Get next message to display
+        try:
+            newmsg = self.all_envelopes[self.current_message_index]
+            # Skip headers
+            if newmsg.get("type") == "header":
+                if self.current_message_index + 1 < len(self.all_envelopes):
+                    newmsg = self.all_envelopes[self.current_message_index + 1]
+                else:
+                    # If we're at the end, go to the previous message
+                    newmsg = self.all_envelopes[self.current_message_index - 1]
+                    self.current_message_index -= 1
+
+            # Show the next message
+            if "id" in newmsg:
+                self.show_message(newmsg["id"])
+        except (IndexError, KeyError):
+            # If no more messages, just reload envelopes
+            self.reload_needed = True
+            self.fetch_envelopes()
 
     def action_open(self) -> None:
         action_open(self)