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)

maildir_gtd/screens/DocumentViewer.py

@@ -82,15 +82,15 @@ class DocumentViewerScreen(Screen):
         """Compose the document viewer screen."""
         yield Container(
             Horizontal(
+                Container(
+                  Button("✕", id="close_button"),
+                  id="button_container"
+                ),
                 Container(
                   Label(f"Viewing: {self.item_name}", id="document_title"),
                   Label(f'[link="{self.web_url}"]Open on Web[/link] | [link="{self.download_url}"]Download File[/link]', id="document_link"),
                 ),
-                Container(
-                  Button("Close", id="close_button"),
-                  id="button_container"
-               ),
-               id="top_container"
+                id="top_container"
             ),
             ScrollableContainer(
                 Markdown("", id="markdown_content"),