feat: add mailbox parameter to download_attachment method (#90)

Previously download_attachment was hardcoded to use INBOX, while other
methods (get_email_count, get_emails_metadata_stream, get_email_body_by_id,
delete_emails) accept a mailbox parameter with INBOX as default.

This change adds consistent mailbox parameter support to download_attachment,
allowing users to download attachments from any mailbox (e.g., All Mail,
Sent, [Gmail]/Sent Mail).

Changes:
- Add mailbox parameter to EmailClient.download_attachment() with default "INBOX"
- Add mailbox parameter to ClassicEmailHandler.download_attachment()
- Add mailbox parameter to abstract EmailHandler.download_attachment()
- Add mailbox parameter to MCP tool download_attachment()
- Add comprehensive tests for mailbox parameter functionality
- Update existing test assertions for new parameter

Author: Wh1isper

mcp_email_server/app.py

@@ -193,6 +193,7 @@ async def download_attachment(
         str, Field(description="The name of the attachment to download (as shown in the attachments list).")
     ],
     save_path: Annotated[str, Field(description="The absolute path where the attachment should be saved.")],
+    mailbox: Annotated[str, Field(description="The mailbox to search in (default: INBOX).")] = "INBOX",
 ) -> AttachmentDownloadResponse:
     settings = get_settings()
     if not settings.enable_attachment_download:
@@ -202,4 +203,4 @@ async def download_attachment(
         raise PermissionError(msg)
 
     handler = dispatch_handler(account_name)
-    return await handler.download_attachment(email_id, attachment_name, save_path)
+    return await handler.download_attachment(email_id, attachment_name, save_path, mailbox)

mcp_email_server/emails/__init__.py

@@ -61,7 +61,17 @@ async def download_attachment(
         email_id: str,
         attachment_name: str,
         save_path: str,
+        mailbox: str = "INBOX",
     ) -> "AttachmentDownloadResponse":
         """
-        Download an email attachment and save it to the specified path
+        Download an email attachment and save it to the specified path.
+
+        Args:
+            email_id: The UID of the email containing the attachment.
+            attachment_name: The filename of the attachment to download.
+            save_path: The local path where the attachment will be saved.
+            mailbox: The mailbox to search in (default: "INBOX").
+
+        Returns:
+            AttachmentDownloadResponse with download result information.
         """

mcp_email_server/emails/classic.py

@@ -443,16 +443,27 @@ async def download_attachment(
         email_id: str,
         attachment_name: str,
         save_path: str,
+        mailbox: str = "INBOX",
     ) -> dict[str, Any]:
-        """Download a specific attachment from an email and save it to disk."""
+        """Download a specific attachment from an email and save it to disk.
+
+        Args:
+            email_id: The UID of the email containing the attachment.
+            attachment_name: The filename of the attachment to download.
+            save_path: The local path where the attachment will be saved.
+            mailbox: The mailbox to search in (default: "INBOX").
+
+        Returns:
+            A dictionary with download result information.
+        """
         imap = self.imap_class(self.email_server.host, self.email_server.port)
         try:
             await imap._client_task
             await imap.wait_hello_from_server()
 
             await imap.login(self.email_server.user_name, self.email_server.password)
             await _send_imap_id(imap)
-            await imap.select(_quote_mailbox("INBOX"))
+            await imap.select(_quote_mailbox(mailbox))
 
             data = await self._fetch_email_with_formats(imap, email_id)
             if not data:
@@ -846,9 +857,20 @@ async def download_attachment(
         email_id: str,
         attachment_name: str,
         save_path: str,
+        mailbox: str = "INBOX",
     ) -> AttachmentDownloadResponse:
-        """Download an email attachment and save it to the specified path."""
-        result = await self.incoming_client.download_attachment(email_id, attachment_name, save_path)
+        """Download an email attachment and save it to the specified path.
+
+        Args:
+            email_id: The UID of the email containing the attachment.
+            attachment_name: The filename of the attachment to download.
+            save_path: The local path where the attachment will be saved.
+            mailbox: The mailbox to search in (default: "INBOX").
+
+        Returns:
+            AttachmentDownloadResponse with download result information.
+        """
+        result = await self.incoming_client.download_attachment(email_id, attachment_name, save_path, mailbox)
         return AttachmentDownloadResponse(
             email_id=result["email_id"],
             attachment_name=result["attachment_name"],

tests/test_classic_handler.py

@@ -285,7 +285,7 @@ async def test_download_attachment(self, classic_handler, tmp_path):
             assert result.size == 1024
             assert result.saved_path == save_path
 
-            mock_download.assert_called_once_with("123", "document.pdf", save_path)
+            mock_download.assert_called_once_with("123", "document.pdf", save_path, "INBOX")
 
     @pytest.mark.asyncio
     async def test_send_email_with_reply_headers(self, classic_handler):

tests/test_email_attachments.py

@@ -210,3 +210,91 @@ async def test_mime_type_detection(self, email_client, tmp_path):
             message_str = str(message)
             for filename in files:
                 assert filename in message_str
+
+
+class TestDownloadAttachmentMailboxParam:
+    """Tests for download_attachment mailbox parameter."""
+
+    @pytest.mark.asyncio
+    async def test_download_attachment_default_mailbox(self, email_client, tmp_path):
+        """Test download_attachment uses INBOX by default."""
+        import asyncio
+
+        save_path = str(tmp_path / "attachment.pdf")
+
+        mock_imap = AsyncMock()
+        mock_imap._client_task = asyncio.Future()
+        mock_imap._client_task.set_result(None)
+        mock_imap.wait_hello_from_server = AsyncMock()
+        mock_imap.login = AsyncMock()
+        mock_imap.select = AsyncMock(return_value=("OK", [b"1"]))
+        mock_imap.logout = AsyncMock()
+
+        # Mock _fetch_email_with_formats to return None (will raise ValueError)
+        with patch.object(email_client, "_fetch_email_with_formats", return_value=None):
+            with patch.object(email_client, "imap_class", return_value=mock_imap):
+                with pytest.raises(ValueError):
+                    await email_client.download_attachment(
+                        email_id="123",
+                        attachment_name="document.pdf",
+                        save_path=save_path,
+                    )
+
+                # Verify select was called with quoted INBOX
+                mock_imap.select.assert_called_once_with('"INBOX"')
+
+    @pytest.mark.asyncio
+    async def test_download_attachment_custom_mailbox(self, email_client, tmp_path):
+        """Test download_attachment with custom mailbox parameter."""
+        import asyncio
+
+        save_path = str(tmp_path / "attachment.pdf")
+
+        mock_imap = AsyncMock()
+        mock_imap._client_task = asyncio.Future()
+        mock_imap._client_task.set_result(None)
+        mock_imap.wait_hello_from_server = AsyncMock()
+        mock_imap.login = AsyncMock()
+        mock_imap.select = AsyncMock(return_value=("OK", [b"1"]))
+        mock_imap.logout = AsyncMock()
+
+        with patch.object(email_client, "_fetch_email_with_formats", return_value=None):
+            with patch.object(email_client, "imap_class", return_value=mock_imap):
+                with pytest.raises(ValueError):
+                    await email_client.download_attachment(
+                        email_id="123",
+                        attachment_name="document.pdf",
+                        save_path=save_path,
+                        mailbox="All Mail",
+                    )
+
+                # Verify select was called with quoted custom mailbox
+                mock_imap.select.assert_called_once_with('"All Mail"')
+
+    @pytest.mark.asyncio
+    async def test_download_attachment_special_folder(self, email_client, tmp_path):
+        """Test download_attachment with special folder like [Gmail]/Sent Mail."""
+        import asyncio
+
+        save_path = str(tmp_path / "attachment.pdf")
+
+        mock_imap = AsyncMock()
+        mock_imap._client_task = asyncio.Future()
+        mock_imap._client_task.set_result(None)
+        mock_imap.wait_hello_from_server = AsyncMock()
+        mock_imap.login = AsyncMock()
+        mock_imap.select = AsyncMock(return_value=("OK", [b"1"]))
+        mock_imap.logout = AsyncMock()
+
+        with patch.object(email_client, "_fetch_email_with_formats", return_value=None):
+            with patch.object(email_client, "imap_class", return_value=mock_imap):
+                with pytest.raises(ValueError):
+                    await email_client.download_attachment(
+                        email_id="123",
+                        attachment_name="document.pdf",
+                        save_path=save_path,
+                        mailbox="[Gmail]/Sent Mail",
+                    )
+
+                # Verify select was called with quoted special folder
+                mock_imap.select.assert_called_once_with('"[Gmail]/Sent Mail"')

tests/test_mcp_tools.py

@@ -480,7 +480,7 @@ async def test_download_attachment_enabled(self):
                 assert result.size == 1024
 
                 mock_handler.download_attachment.assert_called_once_with(
-                    "12345", "document.pdf", "/var/downloads/document.pdf"
+                    "12345", "document.pdf", "/var/downloads/document.pdf", "INBOX"
                 )
 
     @pytest.mark.asyncio