Improve batch fetch with parallel chunks and better code quality

- Parallelize _batch_fetch_dates using asyncio.gather for better performance
- Improve variable naming (t0/t1/t2 -> descriptive names)
- Restore helpful comments for code readability
- Fix type handling in _batch_fetch_headers (accept both bytes and str)
- Add comprehensive tests for batch methods

Co-Authored-By: Paintress <paintress@arcoer.com>

Author: 2 people

mcp_email_server/emails/classic.py

@@ -1,3 +1,4 @@
+import asyncio
 import email.utils
 import mimetypes
 import re
@@ -268,52 +269,74 @@ def _parse_headers(self, email_id: str, raw_headers: bytes) -> dict[str, Any] |
             logger.error(f"Error parsing email headers: {e!s}")
             return None
 
+    async def _fetch_dates_chunk(
+        self,
+        imap: aioimaplib.IMAP4_SSL | aioimaplib.IMAP4,
+        chunk: list[bytes],
+        chunk_num: int,
+        total_chunks: int,
+    ) -> dict[str, datetime]:
+        """Fetch INTERNALDATE for a single chunk of UIDs."""
+        uid_list = ",".join(uid.decode() for uid in chunk)
+        chunk_start = time.perf_counter()
+        _, data = await imap.uid("fetch", uid_list, "(INTERNALDATE)")
+        chunk_elapsed = time.perf_counter() - chunk_start
+
+        chunk_dates: dict[str, datetime] = {}
+        for item in data:
+            if not isinstance(item, bytes) or b"INTERNALDATE" not in item:
+                continue
+            uid_match = re.search(rb"UID (\d+)", item)
+            date_match = re.search(rb'INTERNALDATE "([^"]+)"', item)
+            if uid_match and date_match:
+                uid = uid_match.group(1).decode()
+                date_str = date_match.group(1).decode().strip()
+                chunk_dates[uid] = datetime.strptime(date_str, "%d-%b-%Y %H:%M:%S %z")
+
+        if total_chunks > 1:
+            logger.info(f"Fetched dates chunk {chunk_num}/{total_chunks}: {len(chunk)} UIDs in {chunk_elapsed:.2f}s")
+
+        return chunk_dates
+
     async def _batch_fetch_dates(
         self,
         imap: aioimaplib.IMAP4_SSL | aioimaplib.IMAP4,
         email_ids: list[bytes],
         chunk_size: int = 5000,
     ) -> dict[str, datetime]:
-        """Batch fetch INTERNALDATE for all UIDs in chunks."""
+        """Batch fetch INTERNALDATE for all UIDs in parallel chunks."""
         if not email_ids:
             return {}
 
+        # Split into chunks
+        chunks = [email_ids[i : i + chunk_size] for i in range(0, len(email_ids), chunk_size)]
+        total_chunks = len(chunks)
+
+        # Fetch all chunks in parallel
+        tasks = [
+            self._fetch_dates_chunk(imap, chunk, chunk_num, total_chunks) for chunk_num, chunk in enumerate(chunks, 1)
+        ]
+        results = await asyncio.gather(*tasks)
+
+        # Merge results
         uid_dates: dict[str, datetime] = {}
-        total_chunks = (len(email_ids) + chunk_size - 1) // chunk_size
-        for chunk_num, i in enumerate(range(0, len(email_ids), chunk_size), 1):
-            chunk = email_ids[i : i + chunk_size]
-            uid_list = ",".join(uid.decode() for uid in chunk)
-            chunk_start = time.perf_counter()
-            _, data = await imap.uid("fetch", uid_list, "(INTERNALDATE)")
-            chunk_elapsed = time.perf_counter() - chunk_start
-
-            for item in data:
-                if not isinstance(item, bytes) or b"INTERNALDATE" not in item:
-                    continue
-                uid_match = re.search(rb"UID (\d+)", item)
-                date_match = re.search(rb'INTERNALDATE "([^"]+)"', item)
-                if uid_match and date_match:
-                    uid = uid_match.group(1).decode()
-                    date_str = date_match.group(1).decode().strip()
-                    uid_dates[uid] = datetime.strptime(date_str, "%d-%b-%Y %H:%M:%S %z")
-
-            if total_chunks > 1:
-                logger.info(
-                    f"Fetched dates chunk {chunk_num}/{total_chunks}: {len(chunk)} UIDs in {chunk_elapsed:.2f}s"
-                )
+        for chunk_dates in results:
+            uid_dates.update(chunk_dates)
 
         return uid_dates
 
     async def _batch_fetch_headers(
         self,
         imap: aioimaplib.IMAP4_SSL | aioimaplib.IMAP4,
-        email_ids: list[str],
+        email_ids: list[bytes] | list[str],
     ) -> dict[str, dict[str, Any]]:
         """Batch fetch headers for a list of UIDs."""
         if not email_ids:
             return {}
 
-        uid_list = ",".join(email_ids)
+        # Normalize to list of strings
+        str_ids = [uid.decode() if isinstance(uid, bytes) else uid for uid in email_ids]
+        uid_list = ",".join(str_ids)
         _, data = await imap.uid("fetch", uid_list, "BODY.PEEK[HEADER]")
 
         results: dict[str, dict[str, Any]] = {}
@@ -390,8 +413,11 @@ async def get_emails_metadata_stream(
     ) -> AsyncGenerator[dict[str, Any], None]:
         imap = self.imap_class(self.email_server.host, self.email_server.port)
         try:
+            # Wait for the connection to be established
             await imap._client_task
             await imap.wait_hello_from_server()
+
+            # Login and select mailbox
             await imap.login(self.email_server.user_name, self.email_server.password)
             await _send_imap_id(imap)
             await imap.select(_quote_mailbox(mailbox))
@@ -408,19 +434,21 @@ async def get_emails_metadata_stream(
             )
             logger.info(f"Get metadata: Search criteria: {search_criteria}")
 
+            # Search for messages - use UID SEARCH for better compatibility
             _, messages = await imap.uid_search(*search_criteria)
 
+            # Handle empty or None responses
             if not messages or not messages[0]:
                 logger.warning("No messages returned from search")
                 return
 
             email_ids = messages[0].split()
             logger.info(f"Found {len(email_ids)} email IDs")
 
-            # Phase 1: Batch fetch INTERNALDATE for sorting
-            t0 = time.perf_counter()
+            # Phase 1: Batch fetch INTERNALDATE for sorting (parallel chunks)
+            fetch_dates_start = time.perf_counter()
             uid_dates = await self._batch_fetch_dates(imap, email_ids)
-            t1 = time.perf_counter()
+            fetch_dates_elapsed = time.perf_counter() - fetch_dates_start
 
             # Sort by INTERNALDATE
             sorted_uids = sorted(uid_dates.items(), key=lambda x: x[1], reverse=(order == "desc"))
@@ -430,16 +458,17 @@ async def get_emails_metadata_stream(
             page_uids = [uid for uid, _ in sorted_uids[start : start + page_size]]
 
             if not page_uids:
-                logger.info(f"Phase 1 (dates): {len(uid_dates)} UIDs in {t1 - t0:.2f}s, page {page} empty")
+                logger.info(f"Phase 1 (dates): {len(uid_dates)} UIDs in {fetch_dates_elapsed:.2f}s, page {page} empty")
                 return
 
             # Phase 2: Batch fetch headers for requested page only
+            fetch_headers_start = time.perf_counter()
             metadata_by_uid = await self._batch_fetch_headers(imap, page_uids)
-            t2 = time.perf_counter()
+            fetch_headers_elapsed = time.perf_counter() - fetch_headers_start
 
             logger.info(
-                f"Fetched page {page}: {t1 - t0:.2f}s dates ({len(uid_dates)} UIDs), "
-                f"{t2 - t1:.2f}s headers ({len(page_uids)} UIDs)"
+                f"Fetched page {page}: {fetch_dates_elapsed:.2f}s dates ({len(uid_dates)} UIDs), "
+                f"{fetch_headers_elapsed:.2f}s headers ({len(page_uids)} UIDs)"
             )
 
             # Yield in sorted order

tests/test_classic_handler.py

@@ -364,3 +364,111 @@ async def test_get_emails_content_includes_message_id(self, classic_handler):
 
             # Verify the client method was called correctly
             mock_get_body.assert_called_once_with("123", "INBOX")
+
+
+class TestEmailClientBatchMethods:
+    """Test batch fetch methods for performance optimization."""
+
+    @pytest.fixture
+    def email_client(self, email_settings):
+        return EmailClient(email_settings.incoming)
+
+    def test_parse_headers(self, email_client):
+        """Test _parse_headers method parses email headers correctly."""
+        raw_headers = b"""From: sender@example.com
+To: recipient@example.com
+Cc: cc@example.com
+Subject: Test Subject
+Date: Mon, 20 Jan 2025 10:30:00 +0000
+
+"""
+        result = email_client._parse_headers("123", raw_headers)
+
+        assert result is not None
+        assert result["email_id"] == "123"
+        assert result["subject"] == "Test Subject"
+        assert result["from"] == "sender@example.com"
+        assert "recipient@example.com" in result["to"]
+        assert "cc@example.com" in result["to"]
+        assert result["attachments"] == []
+
+    def test_parse_headers_with_invalid_data(self, email_client):
+        """Test _parse_headers handles malformed headers gracefully."""
+        # Completely broken data that can't be parsed
+        raw_headers = b"\xff\xfe\x00\x00"
+        result = email_client._parse_headers("123", raw_headers)
+
+        # Should return None or a valid dict with fallback values
+        # The implementation catches exceptions and returns None
+        assert result is None or isinstance(result, dict)
+
+    def test_parse_headers_missing_date(self, email_client):
+        """Test _parse_headers handles missing date with fallback."""
+        raw_headers = b"""From: sender@example.com
+To: recipient@example.com
+Subject: No Date Email
+
+"""
+        result = email_client._parse_headers("123", raw_headers)
+
+        assert result is not None
+        assert result["email_id"] == "123"
+        assert result["date"] is not None  # Should have fallback to now()
+
+    @pytest.mark.asyncio
+    async def test_batch_fetch_dates_empty_list(self, email_client):
+        """Test _batch_fetch_dates with empty list returns empty dict."""
+        mock_imap = AsyncMock()
+        result = await email_client._batch_fetch_dates(mock_imap, [])
+
+        assert result == {}
+        mock_imap.uid.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_batch_fetch_headers_empty_list(self, email_client):
+        """Test _batch_fetch_headers with empty list returns empty dict."""
+        mock_imap = AsyncMock()
+        result = await email_client._batch_fetch_headers(mock_imap, [])
+
+        assert result == {}
+        mock_imap.uid.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_batch_fetch_dates_parses_response(self, email_client):
+        """Test _batch_fetch_dates correctly parses IMAP INTERNALDATE response."""
+        mock_imap = AsyncMock()
+        # Simulate IMAP response format for INTERNALDATE
+        mock_imap.uid.return_value = (
+            "OK",
+            [
+                b'1 FETCH (UID 100 INTERNALDATE "20-Jan-2025 10:30:00 +0000")',
+                b'2 FETCH (UID 101 INTERNALDATE "21-Jan-2025 11:00:00 +0000")',
+            ],
+        )
+
+        result = await email_client._batch_fetch_dates(mock_imap, [b"100", b"101"])
+
+        assert "100" in result
+        assert "101" in result
+        assert result["100"].day == 20
+        assert result["101"].day == 21
+
+    @pytest.mark.asyncio
+    async def test_batch_fetch_headers_parses_response(self, email_client):
+        """Test _batch_fetch_headers correctly parses IMAP BODY[HEADER] response."""
+        mock_imap = AsyncMock()
+        # Simulate IMAP response format for BODY[HEADER]
+        mock_imap.uid.return_value = (
+            "OK",
+            [
+                b"1 FETCH (UID 100 BODY[HEADER] {100}",
+                bytearray(b"From: sender@example.com\r\nTo: recipient@example.com\r\nSubject: Test\r\n\r\n"),
+                b")",
+            ],
+        )
+
+        result = await email_client._batch_fetch_headers(mock_imap, ["100"])
+
+        assert "100" in result
+        assert result["100"]["subject"] == "Test"
+        assert result["100"]["from"] == "sender@example.com"