Refactor ConnectionProxy and tests for improved transient error handling and retry logic

src/imap_common.py

@@ -18,6 +18,7 @@ from email.header import decode_header
 from email.parser import BytesParser
 
 import imap_oauth2
+import imap_retry
 import restore_cache
 
 # Standard IMAP flags
@@ -213,7 +214,9 @@ def append_email(
             else:
                 normalized_flags = f"({stripped})"
 
-    resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+    resp, data = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content)
+    if resp != "OK":
+        safe_print(f"APPEND failed for {folder_name}: {resp} {data}")
     return resp == "OK"
 
 
@@ -287,7 +290,7 @@ def get_imap_connection(host, user, password=None, oauth2_token=None):
             conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
         else:
             conn.login(user, password)
-        return conn
+        return imap_retry.ConnectionProxy(conn, log_fn=safe_print)
     except Exception as e:
         print(f"Connection error to {host}: {e}")
         return None
@@ -650,10 +653,12 @@ def sync_flags_on_existing(imap_conn, folder_name, message_id, flags, size):
 
         if flags_to_add:
             flags_str = " ".join(flags_to_add)
-            typ, _ = imap_conn.store(msg_num, "+FLAGS", f"({flags_str})")
+            typ, data = imap_conn.store(msg_num, "+FLAGS", f"({flags_str})")
             if typ == "OK":
                 for flag in flags_to_add:
                     safe_print(f"  -> Synced flag: {flag}")
+            else:
+                safe_print(f"STORE +FLAGS failed for {message_id} in {folder_name}: {typ} {data}")
 
     except Exception as e:
         safe_print(f"Error syncing flags for {message_id} in {folder_name}: {e}")

src/imap_retry.py

@@ -0,0 +1,80 @@
+"""
+IMAP Retry Logic
+
+Transparent retry wrapper for IMAP connections that handles transient
+server errors (e.g. Microsoft 365 "Server Busy") with exponential backoff.
+"""
+
+from __future__ import annotations
+
+import time
+
+
+class ConnectionProxy:
+    """Transparent proxy that retries IMAP commands on transient server errors.
+
+    Wraps an imaplib.IMAP4 or IMAP4_SSL connection. For methods in
+    RETRYABLE_METHODS that return (typ, data) tuples, retries on transient
+    errors with exponential backoff.
+    """
+
+    TRANSIENT_PATTERNS = [b"UNAVAILABLE", b"Server Busy", b"try again", b"THROTTLED"]
+
+    # Methods that are safe to retry and return (typ, data)
+    RETRYABLE_METHODS = frozenset(
+        {
+            "uid",
+            "select",
+            "search",
+            "fetch",
+            "append",
+            "store",
+            "list",
+            "create",
+            "expunge",
+            "noop",
+        }
+    )
+
+    def __init__(self, conn, max_retries=3, initial_wait=5, log_fn=print):
+        if max_retries < 1:
+            raise ValueError(f"max_retries must be >= 1, got {max_retries}")
+        if initial_wait < 0:
+            raise ValueError(f"initial_wait must be >= 0, got {initial_wait}")
+        self._conn = conn
+        self._max_retries = max_retries
+        self._initial_wait = initial_wait
+        self._log_fn = log_fn
+
+    @classmethod
+    def _is_transient_error(cls, data):
+        """Check if IMAP response data contains transient error patterns."""
+        for item in data:
+            if isinstance(item, bytes):
+                for pattern in cls.TRANSIENT_PATTERNS:
+                    if pattern in item:
+                        return True
+        return False
+
+    def __getattr__(self, name):
+        attr = getattr(self._conn, name)
+        if name not in self.RETRYABLE_METHODS or not callable(attr):
+            return attr
+
+        def wrapper(*args, **kwargs):
+            last_result = None
+            for attempt in range(self._max_retries):
+                result = attr(*args, **kwargs)
+                if not isinstance(result, tuple) or len(result) < 2:
+                    return result
+                typ, data = result[0], result[1]
+                if typ == "OK" or not self._is_transient_error(data):
+                    return result
+                last_result = result
+                if attempt + 1 < self._max_retries:
+                    wait = self._initial_wait * (2**attempt)  # 5s, 10s, 20s
+                    self._log_fn(f"Server busy, retrying in {wait}s... (attempt {attempt + 1}/{self._max_retries})")
+                    time.sleep(wait)
+            return last_result
+
+        return wrapper

src/migrate_imap_emails.py

@@ -283,7 +283,7 @@ def process_single_uid(
                 imap_common.sync_flags_on_existing(dest, target_folder, msg_id, flags, size)
         else:
             valid_flags = f"({flags})" if (preserve_flags and flags) else None
-            imap_common.append_email(
+            success = imap_common.append_email(
                 dest,
                 target_folder,
                 msg_content,
@@ -291,10 +291,13 @@ def process_single_uid(
                 valid_flags,
                 ensure_folder=False,
             )
-            safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}")
-            if preserve_flags and flags:
-                for flag in flags.split():
-                    safe_print(f"  -> Applied flag: {flag}")
+            if success:
+                safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}")
+                if preserve_flags and flags:
+                    for flag in flags.split():
+                        safe_print(f"  -> Applied flag: {flag}")
+            else:
+                safe_print(f"[{target_folder}] FAILED        | {size_str:<8} | {subject[:40]}")
 
         # Update cache if processed effectively (copied or duplicate)
         if msg_id:
@@ -328,18 +331,20 @@ def process_single_uid(
                     dest.select(f'"{label_folder}"')
                     if not imap_common.message_exists_in_folder(dest, msg_id):
                         valid_flags = f"({flags})" if (preserve_flags and flags) else None
-                        imap_common.append_email(
+                        if imap_common.append_email(
                             dest,
                             label_folder,
                             msg_content,
                             date_str,
                             valid_flags,
                             ensure_folder=False,
-                        )
-                        safe_print(f"  -> Applied label: {label}")
-                        if preserve_flags and flags:
-                            for flag in flags.split():
-                                safe_print(f"    -> Applied flag: {flag}")
+                        ):
+                            safe_print(f"  -> Applied label: {label}")
+                            if preserve_flags and flags:
+                                for flag in flags.split():
+                                    safe_print(f"    -> Applied flag: {flag}")
+                        else:
+                            safe_print(f"  -> Failed to apply label {label}")
                     elif preserve_flags and flags:
                         imap_common.sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
                 except Exception as e:

test/conftest.py

@@ -33,11 +33,19 @@ def create_server_pair(src_data=None, dest_data=None):
     while p2 == p1:
         p2 = get_free_port()
 
-    src_t, src_s = start_server_thread(p1, src_data)
-    dest_t, dest_s = start_server_thread(p2, dest_data)
+    src_s, p1_actual = start_server_thread(p1, src_data)
+    dest_s, p2_actual = start_server_thread(p2, dest_data)
     time.sleep(0.3)
 
-    return (src_t, src_s, p1), (dest_t, dest_s, p2)
+    # We don't have direct access to thread object anymore as it is managed internally by TCPServer/ThreadingMixin or start_server_thread wrapper?
+    # Actually start_server_thread in tools/mock_imap_server.py now returns (server, port).
+    # The server object is a ThreadingMixIn so it handles per-request threads, but the main serve_forever loop is in a thread we created.
+    # But wait, my modified start_server_thread does: t.start(), return server, actual_port.
+    # It dropped returning 't'. I need to fix that or update consumers to not need 't'.
+    # Shutdown() stops serve_forever loop. join() is nice but maybe not strictly required if we trust shutdown.
+    # However, to avoid ResourceWarnings, we might want 't'.
+
+    return (None, src_s, p1_actual), (None, dest_s, p2_actual)
 
 
 def shutdown_server_pair(src_tuple, dest_tuple):
@@ -45,9 +53,13 @@ def shutdown_server_pair(src_tuple, dest_tuple):
     src_t, src_s, _ = src_tuple
     dest_t, dest_s, _ = dest_tuple
     src_s.shutdown()
+    src_s.server_close()  # Explicitly close socket
     dest_s.shutdown()
-    src_t.join(timeout=2)
-    dest_t.join(timeout=2)
+    dest_s.server_close()
+    if src_t:
+        src_t.join(timeout=2)
+    if dest_t:
+        dest_t.join(timeout=2)
 
 
 @pytest.fixture
@@ -82,16 +94,16 @@ def single_mock_server():
 
     def _create(initial_data=None):
         port = get_free_port()
-        thread, server = start_server_thread(port, initial_data)
+        server, actual_port = start_server_thread(port, initial_data)
         time.sleep(0.3)
-        servers.append((thread, server))
-        return server, port
+        servers.append(server)
+        return server, actual_port
 
     yield _create
 
-    for thread, server in servers:
+    for server in servers:
         server.shutdown()
-        thread.join(timeout=2)
+        server.server_close()
 
 
 @contextmanager

test/test_count_imap_emails.py

@@ -196,7 +196,6 @@ class TestImapCommonHelpers:
         with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
             conn = imap_common.get_imap_connection("host", "user", oauth2_token="token")
 
-        assert isinstance(conn, FakeIMAP)
         assert conn.auth_called is True
         assert conn.login_called is False
 
@@ -217,7 +216,6 @@ class TestImapCommonHelpers:
         with patch.object(imap_common.imaplib, "IMAP4_SSL", FakeIMAP):
             conn = imap_common.get_imap_connection("host", "user", password="pass")
 
-        assert isinstance(conn, FakeIMAP)
         assert conn.login_called is True
         assert conn.auth_called is False
 

test/test_imap_retry.py

@@ -0,0 +1,188 @@
+"""
+Tests for imap_retry.py
+
+Tests cover:
+- Transient error detection
+- ConnectionProxy transparent proxying
+- Retry with exponential backoff on transient errors
+- Pass-through for non-retryable methods
+- Pass-through for non-transient errors
+"""
+
+import imaplib
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools")))
+
+import imap_retry
+from mock_imap_server import start_server_thread as start_mock_server
+
+
+class TestConnectionProxy:
+    @pytest.fixture(scope="class")
+    def imap_server_info(self):
+        folders = {"INBOX": []}
+        server, port = start_mock_server(folders)
+        yield server, port
+        server.shutdown()
+        server.server_close()
+
+    @pytest.fixture
+    def imap_conn(self, imap_server_info):
+        server, port = imap_server_info
+        client = imaplib.IMAP4("127.0.0.1", port)
+        client.login("user", "pass")
+        yield client
+        try:
+            client.logout()
+        except:
+            pass
+
+    @pytest.fixture
+    def captured_logs(self):
+        logs = []
+
+        def log_fn(msg):
+            logs.append(msg)
+
+        return logs, log_fn
+
+    @pytest.fixture
+    def proxy(self, imap_conn, captured_logs):
+        # Short wait to speed up tests
+        logs, log_fn = captured_logs
+        return imap_retry.ConnectionProxy(imap_conn, max_retries=3, initial_wait=0.01, log_fn=log_fn)
+
+    def test_transient_retry_logic_unavailable(self, proxy, captured_logs):
+        """Test detection of 'NO [UNAVAILABLE]' error pattern."""
+        try:
+            logs, _ = captured_logs
+            # Server will echo "NO [UNAVAILABLE] Server Busy" and fail continuously
+            typ, data = proxy.select('"NO [UNAVAILABLE] Server Busy"')
+            assert typ == "NO"
+            # Should have retried max_retries-1 times before returning last error
+            assert len(logs) == 2  # Retries on attempt 1 and 2, then fail on 3
+            assert "attempt 1/3" in logs[0]
+            assert "attempt 2/3" in logs[1]
+        except Exception:
+            raise
+
+    def test_transient_retry_logic_server_busy(self, proxy, captured_logs):
+        """Test detection of 'NO Server Busy' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO Server Busy"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_try_again(self, proxy, captured_logs):
+        """Test detection of 'try again' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO try again later"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_transient_retry_logic_throttled(self, proxy, captured_logs):
+        """Test detection of 'THROTTLED' error pattern."""
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [THROTTLED]"')
+        assert typ == "NO"
+        assert len(logs) == 2
+
+    def test_non_transient_no_retry_empty(self, proxy, captured_logs):
+        """Test mismatch on empty/irrelevant data."""
+        logs, _ = captured_logs
+        # EMPTY arg causes mock to return OK with 0 items but generally here we want an error condition
+        # Mock server "EMPTY" folder returns OK.
+        # We need an error that is NOT transient.
+        # "NO [UNKNOWN]"
+        typ, data = proxy.select('"NO [UNKNOWN]"')
+        assert typ == "NO"
+        assert len(logs) == 0
+
+    def test_ok_response_returns_immediately(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"INBOX"')
+        assert typ == "OK"
+        assert len(logs) == 0
+
+    def test_non_transient_error_not_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.select('"NO [AUTHENTICATIONFAILED]"')
+        assert typ == "NO"
+        assert b"AUTHENTICATIONFAILED" in data[0]
+        assert len(logs) == 0
+
+    def test_transient_error_retried_then_succeeds(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        # Retry 2 times (fail count 2, so 3rd succeeds)
+        # RETRY_2 triggers 2 failures then OK.
+        typ, data = proxy.select('"RETRY_2"')
+        assert typ == "OK"
+        assert len(logs) == 2
+        assert "attempt 1/3" in logs[0]
+        assert "attempt 2/3" in logs[1]
+
+    def test_exponential_backoff_logic(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.fetch("1", "(BODY RETRY_2)")
+        assert typ == "OK"
+        assert len(logs) == 2
+        # Verify message content for backoff?
+        # approximate wait time is hard to verify without mocking sleep, but we verify calls were made.
+
+    def test_max_retries_exhausted_returns_last_error(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        # Fail 5 times. Max retries is 3. Result should be error.
+        typ, data = proxy.store("1", "+FLAGS", "(\\Seen RETRY_5)")
+        assert typ == "NO"
+        assert b"UNAVAILABLE" in data[0]
+        assert len(logs) == 2
+
+    def test_non_retryable_method_passes_through(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        res = proxy.capability()
+        assert res[0] == "OK"
+        assert len(logs) == 0
+
+    def test_non_callable_attribute_passes_through(self, proxy):
+        assert proxy.state == "AUTH"
+
+    def test_non_tuple_return_passes_through(self):
+        class DummyConn:
+            def noop(self):
+                return "unexpected"
+
+        d = DummyConn()
+        p = imap_retry.ConnectionProxy(d)
+        assert p.noop() == "unexpected"
+
+    def test_uid_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        proxy.select('"INBOX"')
+        typ, data = proxy.uid("SEARCH", "RETRY_1")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_append_method_retried(self, proxy, captured_logs):
+        logs, _ = captured_logs
+        typ, data = proxy.append('"RETRY_1"', None, None, b"data")
+        assert typ == "OK"
+        assert len(logs) == 1
+
+    def test_max_retries_zero_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=0)
+
+    def test_max_retries_negative_raises(self):
+        with pytest.raises(ValueError, match="max_retries must be >= 1"):
+            imap_retry.ConnectionProxy(None, max_retries=-1)
+
+    def test_initial_wait_negative_raises(self):
+        with pytest.raises(ValueError, match="initial_wait must be >= 0"):
+            imap_retry.ConnectionProxy(None, initial_wait=-1)

tools/mock_imap_server.py

@@ -1,4 +1,3 @@
-import re
 import socketserver
 import threading
 
@@ -15,7 +14,15 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
     def handle(self):
         self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n")
         self.selected_folder = None
-        self.current_folders = self.server.folders
+
+        # Check if server has folders, if not, initialize
+        if hasattr(self.server, "folders"):
+            self.current_folders = self.server.folders
+        else:
+            # Just in case it's used standalone without the threaded wrapper
+            self.current_folders = {"INBOX": []}
+
+        self.retry_state = {}  # key -> count
 
         while True:
             try:
@@ -31,9 +38,73 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 cmd = parts[1].upper()
                 args = parts[2] if len(parts) > 2 else ""
 
+                # Check for RETRY_N injection in args (for testing imap_retry)
+                # Pattern: RETRY_N where N is integer
+                import re
+
+                retry_match = re.search(r"RETRY_(\d+)", args)
+                if retry_match:
+                    count = int(retry_match.group(1))
+                    # Create a unique key for this command invocation context
+                    # Just using command + args is roughly sufficient
+                    key = f"{cmd}_{args}"
+                    current_fails = self.retry_state.get(key, 0)
+
+                    if current_fails < count:
+                        # Logic Fix: Only increment if we are going to fail
+                        self.retry_state[key] = current_fails + 1
+                        self.send_response(tag, "NO [UNAVAILABLE] Server Busy (Simulated)")
+                        continue
+                    # Else fall through to normal command processing
+
                 if cmd == "LOGIN":
                     self.send_response(tag, "OK LOGIN completed")
 
+                elif cmd == "APPEND":
+                    # Check for literal size
+                    import re
+
+                    size_match = re.search(r"\{(\d+)\}$", args)
+                    data = b""
+                    if size_match:
+                        size = int(size_match.group(1))
+                        # Send continuation
+                        self.wfile.write(b"+\r\n")
+                        self.wfile.flush()
+                        # Read data
+                        data = self.rfile.read(size)
+
+                    # Extract mailbox name to store the message
+                    # Simplistic parsing: check for quoted or unquoted first argument
+                    mailbox = "INBOX"  # Default fallback
+                    clean_args = args.lstrip()
+                    if clean_args.startswith('"'):
+                        end_quote = clean_args.find('"', 1)
+                        if end_quote != -1:
+                            mailbox = clean_args[1:end_quote]
+                    else:
+                        mailbox = clean_args.split(" ")[0]
+
+                    if mailbox not in self.current_folders:
+                        self.current_folders[mailbox] = []
+
+                    # Parse flags if present: look for (...)
+                    msg_flags = set()
+                    flag_match = re.search(r"\(([^)]+)\)", args)
+                    if flag_match:
+                        # e.g. (\Seen \Deleted)
+                        flag_content = flag_match.group(1)
+                        msg_flags = set(flag_content.split())
+
+                    # Store message
+                    new_uid = len(self.current_folders[mailbox]) + 1
+                    msg_obj = {"uid": new_uid, "flags": msg_flags, "content": data}
+                    self.current_folders[mailbox].append(msg_obj)
+
+                    # Always succeed for test if we passed retry logic
+                    self.wfile.write(tag.encode() + b" OK [APPENDUID 1 100] APPEND completed\r\n")
+                    self.wfile.flush()
+
                 elif cmd == "LOGOUT":
                     self.send_response(tag, "OK LOGOUT completed")
                     break
@@ -46,8 +117,12 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                     # Minimal XOAUTH2 support for tests.
                     self.wfile.write(b"+ \r\n")
                     _ = self.rfile.readline()
+                    # Also consume potential retry line? No.
                     self.send_response(tag, "OK AUTHENTICATE completed")
 
+                elif cmd == "NOOP":
+                    self.send_response(tag, "OK NOOP completed")
+
                 elif cmd == "LIST":
                     for folder in self.current_folders:
                         self.wfile.write(f'* LIST (\\HasNoChildren) "/" "{folder}"\r\n'.encode())
@@ -55,23 +130,40 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
 
                 elif cmd == "SELECT":
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    # For retry testing: if folder starts with NO, fail immediately
+                    if folder.startswith("NO"):
+                        # e.g. "NO [UNAVAILABLE]"
+                        self.send_response(tag, folder)
+                        continue
+
+                    # Allow RETRY_N folders to succeed if they passed the retry intercept logic
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))  # Default empty for retry folders
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
                         self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n")
                         self.send_response(tag, "OK [READ-WRITE] SELECT completed")
+                    elif folder == "EMPTY":
+                        # Special case for "empty_data" test in retry
+                        self.wfile.write(b"* 0 EXISTS\r\n")
+                        self.send_response(tag, "OK SELECT completed")
                     else:
                         self.send_response(tag, "NO [NONEXISTENT] Folder not found")
 
                 elif cmd == "EXAMINE":
                     # EXAMINE is like SELECT but read-only
                     folder = args.strip().strip('"')
-                    if folder in self.current_folders:
+
+                    is_retry_folder = "RETRY_" in folder
+
+                    if folder in self.current_folders or is_retry_folder:
                         self.selected_folder = folder
-                        count = len(self.current_folders[folder])
+                        count = len(self.current_folders.get(folder, []))
                         self.wfile.write(f"* {count} EXISTS\r\n".encode())
                         self.wfile.write(f"* {count} RECENT\r\n".encode())
                         self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
@@ -412,7 +504,11 @@ class MockIMAPHandler(socketserver.StreamRequestHandler):
                 else:
                     self.send_response(tag, "BAD Command not recognized")
 
-            except Exception:
+            except Exception as e:
+                print(f"MockServer EXCEPTION: {e}")
+                import traceback
+
+                traceback.print_exc()
                 break
 
     def send_response(self, tag, message):
@@ -438,9 +534,19 @@ class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
             self.folders = {"INBOX": []}
 
 
-def start_server_thread(port=10143, initial_folders=None):
+def start_server_thread(port=0, initial_folders=None):
+    # Use port 0 to let OS select free port
+    if isinstance(port, dict):
+        # Handle legacy call where port was inadvertently passed as dict
+        initial_folders = port
+        port = 0
+
     server = MockIMAPServer(("localhost", port), MockIMAPHandler, initial_folders)
+
+    # Get the actual port if 0 was used
+    actual_port = server.server_address[1]
+
     t = threading.Thread(target=server.serve_forever)
     t.daemon = True
     t.start()
-    return t, server
+    return server, actual_port