Add imap_session module for centralized OAuth2 refresh (#21)

* Add .claude to gitignore

* Add imap_session module for centralized OAuth2 refresh

Create new imap_session.py module that combines imap_common and
imap_oauth2 for unified session management:

- ensure_connection(): Refreshes OAuth2 token and ensures connection health
- ensure_folder_session(): Ensures connection AND folder selection

Updated backup, migrate, and restore scripts to use the new module,
simplifying the proactive token refresh pattern in process_batch
functions.

* Move is_token_expired_error to imap_oauth2 module

Relocates the OAuth2 token expiration detection function to the oauth2-specific module where it logically belongs.

* Add retry logic for auth errors during email processing

Extract single-UID processing into dedicated functions with retry capability
for recoverable authentication failures. Broadens auth error detection beyond
token expiration to include session loss and authentication failures.

* Cache Microsoft tenant discovery results

Avoids repeated OpenID Connect discovery requests for the same email domain
during OAuth2 authentication flows. Normalizes domain to lowercase for
consistent cache keys.

* Add tests for is_token_expired_error and is_auth_error

Covers token expiration, session invalidation, and authentication failure
detection with case-insensitive matching verification. Ensures non-auth
errors like connection and timeout issues are correctly excluded.

* Add tests for imap_session module

Covers ensure_connection and ensure_folder_session functions including
OAuth2 token refresh, folder selection, and error handling scenarios.
Verifies connection change detection triggers folder re-selection.

* Add tests for auth error re-raising in get_message_ids_in_folder

Verifies that authentication errors are propagated to callers for
reconnection handling while non-auth errors return empty results.
Clarifies test names to distinguish error handling behaviors.

.gitignore

@@ -7,3 +7,4 @@ env/
 venv/
 .env
 .coverage
+.claude

src/backup_imap_emails.py

@@ -58,6 +58,7 @@ import threading
 
 import imap_common
 import imap_oauth2
+import imap_session
 
 # Defaults
 MAX_WORKERS = 10
@@ -76,13 +77,6 @@ def safe_print(message):
         print(f"[{short_name}] {message}")
 
 
-def ensure_connection(conn, conf):
-    """Refresh OAuth2 token if needed and ensure connection is healthy."""
-    if conf.get("oauth2"):
-        imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
-    return imap_common.ensure_connection_from_conf(conn, conf)
-
-
 def get_thread_connection(src_conf):
     # Backwards-compatible: tests and some internal call sites may still pass
     # (host, user, password) tuples. Normalize to the dict form used by
@@ -90,10 +84,74 @@ def get_thread_connection(src_conf):
     if isinstance(src_conf, tuple) and len(src_conf) == 3:
         src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
 
-    thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf)
+    thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf)
     return thread_local.src
 
 
+def process_single_uid(src, uid, folder_name, local_folder_path):
+    """
+    Fetch and save a single email by UID.
+
+    Args:
+        src: IMAP connection
+        uid: Message UID to fetch
+        folder_name: Current folder name (for logging)
+        local_folder_path: Local directory to save email
+
+    Returns:
+        Tuple of (success: bool, connection):
+        - (True, src): UID processed successfully (or skipped)
+        - (False, src): Auth error occurred, caller should retry after reconnect
+    """
+    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
+
+    try:
+        resp, data = src.uid("fetch", uid, "(RFC822)")
+        if resp != "OK" or not data or data[0] is None:
+            safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
+            return (True, src)  # Don't retry, move on
+
+        raw_email = None
+        for item in data:
+            if isinstance(item, tuple):
+                raw_email = item[1]
+                break
+
+        # Derive Subject for filename from the already-fetched message bytes.
+        _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
+        if not subject:
+            clean_subject = "No Subject"
+        else:
+            clean_subject = imap_common.sanitize_filename(subject)
+            clean_subject = clean_subject[:100]
+
+        filename = f"{uid_str}_{clean_subject}.eml"
+        full_path = os.path.join(local_folder_path, filename)
+
+        if os.path.exists(full_path):
+            return (True, src)  # Already exists
+
+        if raw_email:
+            try:
+                with open(full_path, "wb") as f:
+                    f.write(raw_email)
+                safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
+            except OSError as e:
+                safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
+        else:
+            safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
+
+        return (True, src)
+
+    except Exception as e:
+        if imap_oauth2.is_auth_error(e):
+            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
+            return (False, src)  # Signal retry needed
+        else:
+            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
+            return (True, src)  # Don't retry other errors
+
+
 def process_batch(uids, folder_name, src_conf, local_folder_path):
     if isinstance(src_conf, tuple) and len(src_conf) == 3:
         src_conf = {"host": src_conf[0], "user": src_conf[1], "password": src_conf[2]}
@@ -110,52 +168,24 @@ def process_batch(uids, folder_name, src_conf, local_folder_path):
         return
 
     for uid in uids:
-        try:
+        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-            # Helper to handle byte UIDs
+        max_retries = 2
-            uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
 
-            # Fetch Full Content
+        for attempt in range(max_retries):
-            # RFC822 gets the whole message including headers and attachments
+            src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
-            resp, data = src.uid("fetch", uid, "(RFC822)")
+            thread_local.src = src
-            if resp != "OK" or not data or data[0] is None:
+            if not ok:
-                safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
+                safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
-                continue
+                return
 
-            raw_email = None
+            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
-            for item in data:
+            thread_local.src = src
-                if isinstance(item, tuple):
-                    raw_email = item[1]
-                    break
 
-            # Derive Subject for filename from the already-fetched message bytes.
+            if success:
-            _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
+                break
-            if not subject:
+            if attempt < max_retries - 1:
-                clean_subject = "No Subject"
+                src = None
-            else:
+                thread_local.src = None
-                clean_subject = imap_common.sanitize_filename(subject)
-                # Limit subject len to avoid FS path length limits (max 255 usually)
-                clean_subject = clean_subject[:100]
-
-            filename = f"{uid_str}_{clean_subject}.eml"
-            full_path = os.path.join(local_folder_path, filename)
-
-            # Double check existence (in case incremental UID checks missed it or naming collision).
-            if os.path.exists(full_path):
-                continue
-
-            if raw_email:
-                try:
-                    with open(full_path, "wb") as f:
-                        f.write(raw_email)
-                    safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
-                except OSError as e:
-                    # Valid filename might still fail if path too long
-                    safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
-            else:
-                safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
-
-        except Exception as e:
-            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
 
 
 def get_existing_uids(local_path):
@@ -210,11 +240,21 @@ def is_gmail_label_folder(folder_name):
 PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
 
 
-def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
+def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
     """
-    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
-    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+    with OAuth2 session management.
-    Optional progress_callback(current, total) for progress reporting.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_info, imap_conn) where message_info is:
+        { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+        The returned imap_conn may be different if reconnection occurred.
     """
     message_info = {}
 
@@ -222,16 +262,16 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
         imap_conn.select(f'"{folder_name}"', readonly=True)
     except Exception as e:
         safe_print(f"Could not select folder {folder_name}: {e}")
-        return message_info
+        return (message_info, imap_conn)
 
     try:
         resp, data = imap_conn.uid("search", None, "ALL")
         if resp != "OK" or not data or not data[0]:
-            return message_info
+            return (message_info, imap_conn)
 
         uids = data[0].split()
         if not uids:
-            return message_info
+            return (message_info, imap_conn)
 
         total_uids = len(uids)
 
@@ -245,6 +285,13 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
             if progress_callback:
                 progress_callback(min(i + batch_size, total_uids), total_uids)
 
+            # Proactively refresh token and ensure folder is selected
+            if src_conf:
+                imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
+                if not ok:
+                    safe_print(f"ERROR: Connection/folder lost in {folder_name}")
+                    break
+
             try:
                 # Fetch both Message-ID header and FLAGS
                 resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
@@ -271,19 +318,44 @@ def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
                             message_info[msg_id] = {"flags": flags}
             except Exception as e:
                 safe_print(f"Error fetching batch in {folder_name}: {e}")
-                # Try to keep connection alive
-                try:
-                    imap_conn.noop()
-                except Exception:
-                    pass
                 continue
 
     except Exception as e:
         safe_print(f"Error searching folder {folder_name}: {e}")
 
+    return (message_info, imap_conn)
+
+
+def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
+    """
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
+
+    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
     return message_info
 
 
+def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
+    """
+    Returns a set of Message-IDs for all emails in a given folder,
+    with OAuth2 session management.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
+        The returned imap_conn may be different if reconnection occurred.
+    """
+    info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
+    return (set(info.keys()), imap_conn)
+
+
 def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
     """
     Returns a set of Message-IDs for all emails in a given folder.
@@ -294,7 +366,7 @@ def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
     return set(info.keys())
 
 
-def build_labels_manifest(imap_conn, local_path):
+def build_labels_manifest(imap_conn, local_path, src_conf=None):
     """
     Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
     Scans all folders (labels) in the account and records which Message-IDs
@@ -302,6 +374,7 @@ def build_labels_manifest(imap_conn, local_path):
 
     Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
     Saves the manifest to labels_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
     """
     import time
 
@@ -329,7 +402,9 @@ def build_labels_manifest(imap_conn, local_path):
             flush=True,
         )
 
-    all_mail_info = get_message_info_in_folder(imap_conn, imap_common.GMAIL_ALL_MAIL, all_mail_progress_cb)
+    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
+        imap_conn, imap_common.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
+    )
     print()  # New line after progress
 
     # Initialize manifest with flags from All Mail
@@ -369,7 +444,7 @@ def build_labels_manifest(imap_conn, local_path):
             )
 
         safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        message_ids = get_message_ids_in_folder(imap_conn, folder_name, progress_cb)
+        message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
         print()  # New line after progress
 
         # Keep connection alive between folders
@@ -422,13 +497,14 @@ def build_labels_manifest(imap_conn, local_path):
     return manifest
 
 
-def build_flags_manifest(imap_conn, local_path, folders_to_scan=None):
+def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
     """
     Builds a manifest mapping Message-IDs to their IMAP flags.
     For non-Gmail servers, scans specified folders (or all folders if not specified).
 
     Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
     Saves the manifest to flags_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
     """
     import time
 
@@ -467,7 +543,7 @@ def build_flags_manifest(imap_conn, local_path, folders_to_scan=None):
             )
 
         safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        folder_info = get_message_info_in_folder(imap_conn, folder_name, progress_cb)
+        folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
         print()  # New line after progress
 
         # Merge into manifest (keep first occurrence of flags)
@@ -829,7 +905,7 @@ def main():
         if args.preserve_labels or args.manifest_only or args.gmail_mode:
             print("Building Gmail labels manifest...")
             print("This scans all folders to map Message-IDs to labels and flags.\n")
-            build_labels_manifest(src, local_path)
+            build_labels_manifest(src, local_path, src_conf)
             print("")  # Blank line after manifest building
         # Build flags-only manifest for non-Gmail servers
         elif args.preserve_flags:
@@ -837,12 +913,15 @@ def main():
             print("This scans folders to capture read/starred/etc status.\n")
             # Get folders to scan
             folders_to_scan = [args.folder] if args.folder else None
-            build_flags_manifest(src, local_path, folders_to_scan)
+            build_flags_manifest(src, local_path, folders_to_scan, src_conf)
             print("")  # Blank line after manifest building
 
         # If manifest-only mode, we're done
         if args.manifest_only:
-            src.logout()
+            try:
+                src.logout()
+            except Exception:
+                pass  # Connection may already be closed
             manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
             print("\nManifest-only mode complete.")
             print(f"Labels manifest saved to: {manifest_path}")
@@ -856,15 +935,23 @@ def main():
         elif args.folder:
             backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
         else:
+            # Reconnect after potentially long manifest building
+            src = imap_session.ensure_connection(src, src_conf)
+            if not src:
+                print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
+                sys.exit(0)
             folders = imap_common.list_selectable_folders(src)
             for name in folders:
-                src = ensure_connection(src, src_conf)
+                src = imap_session.ensure_connection(src, src_conf)
                 if not src:
                     print("Fatal: Could not reconnect to IMAP server. Aborting.")
                     sys.exit(1)
                 backup_folder(src, name, local_path, src_conf, args.dest_delete)
 
-        src.logout()
+        try:
+            src.logout()
+        except Exception:
+            pass  # Connection may already be closed
         print("\nBackup completed successfully.")
 
         if args.preserve_labels or args.gmail_mode:

src/imap_common.py

@@ -14,6 +14,8 @@ from email import policy
 from email.header import decode_header
 from email.parser import BytesParser
 
+import imap_oauth2
+
 # Standard IMAP flags
 FLAG_SEEN = "\\Seen"
 FLAG_ANSWERED = "\\Answered"
@@ -363,7 +365,10 @@ def get_message_ids_in_folder(imap_conn):
         resp, data = imap_conn.uid("search", None, "ALL")
         if resp != "OK" or not data[0].strip():
             return {}
-    except Exception:
+    except Exception as e:
+        # Re-raise token expiration errors so callers can handle reconnection
+        if imap_oauth2.is_auth_error(e):
+            raise
         return {}
 
     uids = data[0].split()
@@ -403,7 +408,10 @@ def get_uid_to_message_id_map(imap_conn, uids):
                     mid = extract_message_id(item[1])
                     if mid:
                         uid_to_msgid[uid] = mid
-        except Exception:
+        except Exception as e:
+            # Re-raise token expiration errors so callers can handle reconnection
+            if imap_oauth2.is_auth_error(e):
+                raise
             continue
 
     return uid_to_msgid

src/imap_oauth2.py

@@ -23,6 +23,7 @@ import urllib.parse
 # Module-level caches for OAuth2 token refresh
 _msal_app_cache = {}  # (client_id, tenant_id) -> PublicClientApplication
 _google_creds_cache = {}  # (client_id, client_secret) -> credentials
+_tenant_cache = {}  # domain -> tenant_id
 _token_refresh_lock = threading.Lock()
 
 
@@ -46,6 +47,35 @@ def _fetch_json_https(host, path, timeout=10):
     return json.loads(body.decode("utf-8"))
 
 
+def is_token_expired_error(error):
+    """
+    Check if an exception indicates OAuth2 token expiration.
+
+    Args:
+        error: The exception to check
+
+    Returns:
+        True if the error indicates token expiration, False otherwise
+    """
+    error_str = str(error).lower()
+    return "accesstokenexpired" in error_str or "session invalidated" in error_str
+
+
+def is_auth_error(error):
+    """
+    Check if an exception indicates an authentication failure that may be recoverable
+    by reconnecting (token expiration, session loss, etc.).
+
+    Args:
+        error: The exception to check
+
+    Returns:
+        True if the error indicates an auth failure, False otherwise
+    """
+    error_str = str(error).lower()
+    return is_token_expired_error(error) or "not authenticated" in error_str or "authentication failed" in error_str
+
+
 def detect_oauth2_provider(host):
     """
     Detects the OAuth2 provider from the IMAP host.
@@ -63,13 +93,18 @@ def discover_microsoft_tenant(email):
     """
     Auto-discovers the Microsoft tenant ID from an email address domain.
     Uses the OpenID Connect discovery endpoint (no authentication required).
+    Results are cached per domain to avoid repeated network requests.
     Returns the tenant ID string or None if discovery fails.
     """
-    domain = email.split("@")[-1].strip()
+    domain = email.split("@")[-1].strip().lower()
     if not domain:
         print("Error: Could not discover Microsoft tenant: missing email domain")
         return None
 
+    # Return cached tenant if available
+    if domain in _tenant_cache:
+        return _tenant_cache[domain]
+
     domain_quoted = urllib.parse.quote(domain, safe=".-")
     path = f"/{domain_quoted}/.well-known/openid-configuration"
 
@@ -82,7 +117,9 @@ def discover_microsoft_tenant(email):
     issuer = data.get("issuer", "")
     match = re.search(r"/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})", issuer)
     if match:
-        return match.group(1)
+        tenant_id = match.group(1)
+        _tenant_cache[domain] = tenant_id
+        return tenant_id
 
     print(f"Error: Could not extract tenant ID from issuer: {issuer}")
     return None

src/imap_session.py

@@ -0,0 +1,62 @@
+"""
+IMAP Session Management
+
+Connection and session management for IMAP operations with OAuth2 support.
+Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh).
+"""
+
+import imap_common
+import imap_oauth2
+
+
+def ensure_connection(conn, conf):
+    """
+    Refresh OAuth2 token if needed and ensure connection is healthy.
+
+    Args:
+        conn: Existing IMAP connection or None
+        conf: Connection config dict with keys:
+            - host, user, password, oauth2_token: connection credentials
+            - oauth2: dict with provider, client_id, email, client_secret (optional)
+
+    Returns:
+        Healthy IMAP connection, or None if connection failed.
+        May return a different connection object if reconnection was needed.
+    """
+    if conf.get("oauth2"):
+        imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
+    return imap_common.ensure_connection_from_conf(conn, conf)
+
+
+def ensure_folder_session(conn, conf, folder_name, readonly=True):
+    """
+    Ensure connection is healthy and folder is selected.
+
+    Proactively refreshes OAuth2 token if needed. If the connection was
+    refreshed (new connection object), reselects the folder.
+
+    Args:
+        conn: Existing IMAP connection or None
+        conf: Connection config dict (see ensure_connection)
+        folder_name: IMAP folder to select
+        readonly: Whether to select folder as readonly (default True)
+
+    Returns:
+        Tuple of (connection, success: bool)
+        - On success: (connection, True) - folder is selected
+        - On failure: (connection or None, False)
+    """
+    old_conn = conn
+    new_conn = ensure_connection(conn, conf)
+
+    if not new_conn:
+        return None, False
+
+    # If connection changed (token refreshed), need to reselect folder
+    if new_conn is not old_conn:
+        try:
+            new_conn.select(f'"{folder_name}"', readonly=readonly)
+        except Exception:
+            return new_conn, False
+
+    return new_conn, True

src/migrate_imap_emails.py

@@ -117,6 +117,7 @@ import threading
 
 import imap_common
 import imap_oauth2
+import imap_session
 
 # Configuration defaults
 DELETE_FROM_SOURCE_DEFAULT = False
@@ -288,19 +289,167 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg
     return deleted_count
 
 
-def ensure_connection(conn, conf):
-    """Refresh OAuth2 token if needed and ensure connection is healthy."""
-    if conf.get("oauth2"):
-        imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
-    return imap_common.ensure_connection_from_conf(conn, conf)
-
-
 def get_thread_connections(src_conf, dest_conf):
-    thread_local.src = ensure_connection(getattr(thread_local, "src", None), src_conf)
+    thread_local.src = imap_session.ensure_connection(getattr(thread_local, "src", None), src_conf)
-    thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf)
+    thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf)
     return thread_local.src, thread_local.dest
 
 
+def process_single_uid(
+    src,
+    dest,
+    uid,
+    folder_name,
+    delete_from_source,
+    trash_folder,
+    preserve_flags,
+    gmail_mode,
+    label_index,
+    check_duplicate,
+):
+    """
+    Migrate a single email by UID.
+
+    Returns:
+        Tuple of (success, src, dest, deleted):
+        - success=True: UID processed (copied, skipped, or non-auth error)
+        - success=False: Auth error, caller should retry after reconnect
+        - deleted: 1 if message was marked for deletion, 0 otherwise
+    """
+    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
+
+    try:
+        resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])")
+        if resp != "OK":
+            safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}")
+            return (True, src, dest, 0)
+
+        msg_content = None
+        flags = None
+        date_str = None
+
+        for item in data:
+            if isinstance(item, tuple):
+                msg_content = item[1]
+                meta = item[0].decode("utf-8", errors="ignore")
+                flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta)
+                if flags_match:
+                    flags = filter_preservable_flags(flags_match.group(1))
+                date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta)
+                if date_match:
+                    date_str = f'"{date_match.group(1)}"'
+
+        if not msg_content:
+            return (True, src, dest, 0)
+
+        size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0
+        size_str = f"{size / 1024:.1f}KB" if size else "0KB"
+        msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content)
+
+        if not msg_id:
+            msg_id = imap_common.parse_message_id_from_bytes(msg_content)
+
+        # Determine target folder and labels for Gmail mode
+        apply_labels = gmail_mode
+        labels = []
+        if apply_labels and msg_id and label_index is not None:
+            labels = sorted(label_index.get(msg_id, set()))
+
+        if apply_labels:
+            skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH}
+            target_folder = None
+            remaining_labels = []
+
+            for label in labels:
+                label_folder = label_to_folder(label)
+                if label_folder in skip_folders:
+                    continue
+                if target_folder is None:
+                    target_folder = label_folder
+                else:
+                    remaining_labels.append(label)
+
+            if target_folder is None:
+                target_folder = imap_common.FOLDER_RESTORED_UNLABELED
+                remaining_labels = []
+        else:
+            target_folder = folder_name
+
+        imap_common.ensure_folder_exists(dest, target_folder)
+        dest.select(f'"{target_folder}"')
+
+        is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id))
+
+        if is_duplicate:
+            safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
+            if preserve_flags and flags and msg_id:
+                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(
+                dest,
+                target_folder,
+                msg_content,
+                date_str,
+                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}")
+
+        # Apply remaining Gmail labels
+        if apply_labels and remaining_labels and msg_id:
+            for label in remaining_labels:
+                label_folder = label_to_folder(label)
+                if label_folder == target_folder:
+                    continue
+                if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH):
+                    continue
+                try:
+                    imap_common.ensure_folder_exists(dest, label_folder)
+                    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(
+                            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}")
+                    elif preserve_flags and flags:
+                        sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
+                except Exception as e:
+                    safe_print(f"  -> Error applying label {label}: {e}")
+
+        deleted = 0
+        if delete_from_source:
+            if trash_folder and folder_name != trash_folder:
+                try:
+                    src.uid("copy", uid, f'"{trash_folder}"')
+                except Exception:
+                    pass
+            src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL)
+            deleted = 1
+
+        return (True, src, dest, deleted)
+
+    except Exception as e:
+        if imap_oauth2.is_auth_error(e):
+            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
+            return (False, src, dest, 0)
+        else:
+            safe_print(f"[{folder_name}] ERROR Exec | UID {uid_str}: {e}")
+            return (True, src, dest, 0)
+
+
 def process_batch(
     uids,
     folder_name,
@@ -318,14 +467,12 @@ def process_batch(
         safe_print("Error: Could not establish connections in worker thread.")
         return
 
-    # Select source folder
     try:
         src.select(f'"{folder_name}"', readonly=False)
     except Exception as e:
         safe_print(f"Error selecting folder {folder_name} in worker: {e}")
         return
 
-    # In non-Gmail-mode, we keep a selected destination folder for efficiency
     if not gmail_mode:
         try:
             imap_common.ensure_folder_exists(dest, folder_name)
@@ -335,134 +482,54 @@ def process_batch(
             return
 
     deleted_count = 0
+
     for uid in uids:
-        try:
+        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-            # Fetch full message (needed to copy and/or apply labels)
+        max_retries = 2
-            resp, data = src.uid("fetch", uid, "(FLAGS INTERNALDATE BODY.PEEK[])")
-            if resp != "OK":
-                uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-                safe_print(f"[{folder_name}] ERROR Fetch | UID {uid_str}")
-                continue
 
-            msg_content = None
+        for attempt in range(max_retries):
-            flags = None
+            src, src_ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=False)
-            date_str = None
+            thread_local.src = src
+            if not src_ok:
+                safe_print(f"[{folder_name}] ERROR: Source connection/folder lost for UID {uid_str}")
+                return
 
-            for item in data:
+            if not gmail_mode:
-                if isinstance(item, tuple):
+                dest, dest_ok = imap_session.ensure_folder_session(dest, dest_conf, folder_name, readonly=False)
-                    msg_content = item[1]
+                thread_local.dest = dest
-                    meta = item[0].decode("utf-8", errors="ignore")
+                if not dest_ok:
-                    flags_match = re.search(r"FLAGS\s+\((.*?)\)", meta)
+                    safe_print(f"[{folder_name}] ERROR: Dest connection/folder lost for UID {uid_str}")
-                    if flags_match:
+                    return
-                        flags = filter_preservable_flags(flags_match.group(1))
-                    date_match = re.search(r"INTERNALDATE\s+\"(.*?)\"", meta)
-                    if date_match:
-                        date_str = f'"{date_match.group(1)}"'
-
-            if not msg_content:
-                continue
-
-            # Compute size and parse headers from the already-fetched message bytes.
-            size = len(msg_content) if isinstance(msg_content, (bytes, bytearray)) else 0
-            size_str = f"{size / 1024:.1f}KB" if size else "0KB"
-            msg_id, subject = imap_common.parse_message_id_and_subject_from_bytes(msg_content)
-
-            # Prefer Message-ID from body if missing from header parse
-            if not msg_id:
-                msg_id = imap_common.parse_message_id_from_bytes(msg_content)
-
-            # Determine target folder and labels for Gmail mode
-            apply_labels = gmail_mode
-            labels = []
-            if apply_labels and msg_id and label_index is not None:
-                labels = sorted(label_index.get(msg_id, set()))
-
-            if apply_labels:
-                skip_folders = {imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH}
-                target_folder = None
-                remaining_labels = []
-
-                for label in labels:
-                    label_folder = label_to_folder(label)
-                    if label_folder in skip_folders:
-                        continue
-                    if target_folder is None:
-                        target_folder = label_folder
-                    else:
-                        remaining_labels.append(label)
-
-                if target_folder is None:
-                    target_folder = imap_common.FOLDER_RESTORED_UNLABELED
-                    remaining_labels = []
             else:
-                target_folder = folder_name
+                dest = imap_session.ensure_connection(dest, dest_conf)
-                remaining_labels = []
+                thread_local.dest = dest
+                if not dest:
+                    safe_print(f"[{folder_name}] ERROR: Dest connection lost for UID {uid_str}")
+                    return
 
-            imap_common.ensure_folder_exists(dest, target_folder)
+            success, src, dest, deleted = process_single_uid(
-            dest.select(f'"{target_folder}"')
+                src,
+                dest,
+                uid,
+                folder_name,
+                delete_from_source,
+                trash_folder,
+                preserve_flags,
+                gmail_mode,
+                label_index,
+                check_duplicate,
+            )
+            thread_local.src = src
+            thread_local.dest = dest
+            deleted_count += deleted
 
-            is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id))
+            if success:
-
+                break
-            if is_duplicate:
+            if attempt < max_retries - 1:
-                safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
+                src = None
-                if preserve_flags and flags and msg_id:
+                dest = None
-                    sync_flags_on_existing(dest, target_folder, msg_id, flags, size)
+                thread_local.src = None
-            else:
+                thread_local.dest = None
-                valid_flags = f"({flags})" if (preserve_flags and flags) else None
-                imap_common.append_email(
-                    dest,
-                    target_folder,
-                    msg_content,
-                    date_str,
-                    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}")
-
-            # Apply remaining Gmail labels (always, whether copied or skipped)
-            if apply_labels and remaining_labels and msg_id:
-                for label in remaining_labels:
-                    label_folder = label_to_folder(label)
-                    if label_folder == target_folder:
-                        continue
-                    if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH):
-                        continue
-                    try:
-                        imap_common.ensure_folder_exists(dest, label_folder)
-                        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(
-                                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}")
-                        elif preserve_flags and flags:
-                            sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
-                    except Exception as e:
-                        safe_print(f"  -> Error applying label {label}: {e}")
-
-            if delete_from_source:
-                if trash_folder and folder_name != trash_folder:
-                    try:
-                        src.uid("copy", uid, f'"{trash_folder}"')
-                    except Exception:
-                        pass
-                src.uid(imap_common.CMD_STORE, uid, imap_common.OP_ADD_FLAGS, imap_common.FLAG_DELETED_LITERAL)
-                deleted_count += 1
-
-        except Exception as e:
-            safe_print(f"[{folder_name}] ERROR Exec | UID {uid}: {e}")
 
     if delete_from_source and deleted_count > 0:
         try:
@@ -961,11 +1028,11 @@ def main():
                         safe_print(f"Skipping migration of Trash folder '{name}' (preventing circular migration).")
                         continue
 
-                    src_main = ensure_connection(src_main, src_conf)
+                    src_main = imap_session.ensure_connection(src_main, src_conf)
                     if not src_main:
                         safe_print("Fatal: Could not reconnect to source IMAP server. Aborting.")
                         sys.exit(1)
-                    dest_main = ensure_connection(dest_main, dest_conf)
+                    dest_main = imap_session.ensure_connection(dest_main, dest_conf)
                     if not dest_main:
                         safe_print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
                         sys.exit(1)

src/restore_imap_emails.py

@@ -61,6 +61,7 @@ from typing import Optional
 
 import imap_common
 import imap_oauth2
+import imap_session
 import restore_cache
 
 
@@ -88,16 +89,9 @@ def safe_print(message):
         print(f"[{short_name}] {message}")
 
 
-def ensure_connection(conn, conf):
-    """Refresh OAuth2 token if needed and ensure connection is healthy."""
-    if conf.get("oauth2"):
-        imap_oauth2.refresh_oauth2_token(conf, conf.get("oauth2_token"))
-    return imap_common.ensure_connection_from_conf(conn, conf)
-
-
 def get_thread_connection(dest_conf):
     """Get or create a thread-local IMAP connection."""
-    thread_local.dest = ensure_connection(getattr(thread_local, "dest", None), dest_conf)
+    thread_local.dest = imap_session.ensure_connection(getattr(thread_local, "dest", None), dest_conf)
     return thread_local.dest
 
 
@@ -395,10 +389,16 @@ def process_restore_batch(
     gmail_mode = folder_name == "__GMAIL_MODE__"
 
     for file_path, filename in eml_files:
+        # Proactively refresh token if needed
+        dest = get_thread_connection(dest_conf)
+        if not dest:
+            safe_print(f"ERROR: Connection lost for {filename}")
+            return
+
         try:
             message_id, date_str, raw_content, subject = parse_eml_file(file_path)
             if raw_content is None:
-                continue
+                continue  # No content, skip to next file
 
             size = len(raw_content)
             size_str = f"{size / 1024:.1f}KB"
@@ -472,7 +472,7 @@ def process_restore_batch(
                 and message_id in existing_dest_msg_ids
             ):
                 safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
-                continue
+                continue  # Skip to next file
 
             # Upload to target folder.
             # Keep server-side duplicate checks enabled to avoid creating duplicates for emails
@@ -815,7 +815,7 @@ def restore_folder(
     # Delete orphan emails from destination if enabled
     if dest_delete and local_msg_ids is not None:
         safe_print("Syncing destination: removing emails not in local backup...")
-        dest = ensure_connection(None, dest_conf)
+        dest = imap_session.ensure_connection(None, dest_conf)
         if dest:
             delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids)
             dest.logout()
@@ -1192,7 +1192,7 @@ def main():
                     continue
 
                 # Proactively refresh OAuth2 token and ensure connection is healthy between folders
-                dest = ensure_connection(dest, dest_conf)
+                dest = imap_session.ensure_connection(dest, dest_conf)
                 if not dest:
                     print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
                     sys.exit(1)

test/test_backup_imap_emails.py

@@ -551,16 +551,16 @@ class TestGmailLabelsPreservation:
             "<msg2@test.com>": {"flags": []},
         }
 
-        def mock_get_message_ids(conn, folder, progress_cb=None):
+        def mock_get_message_info_with_conf(conn, folder, src_conf, progress_cb=None):
-            return folder_data.get(folder, set())
-
-        def mock_get_message_info(conn, folder, progress_cb=None):
             if folder == "[Gmail]/All Mail":
-                return all_mail_info
+                return (all_mail_info, conn)
-            return {}
+            return ({}, conn)
 
-        monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder", mock_get_message_ids)
+        def mock_get_message_ids_with_conf(conn, folder, src_conf, progress_cb=None):
-        monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder", mock_get_message_info)
+            return (folder_data.get(folder, set()), conn)
+
+        monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder_with_conf", mock_get_message_ids_with_conf)
+        monkeypatch.setattr(backup_imap_emails, "get_message_info_in_folder_with_conf", mock_get_message_info_with_conf)
 
         result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))
 

test/test_imap_common.py

@@ -16,6 +16,8 @@ import os
 import sys
 from unittest.mock import Mock
 
+import pytest
+
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
 import imap_common
@@ -653,14 +655,22 @@ class TestGetMessageIdsInFolder:
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
-    def test_search_exception(self):
+    def test_search_non_auth_exception_returns_empty(self):
-        """Test returns empty dict when search raises exception."""
+        """Test returns empty dict when search raises non-auth exception."""
         mock_conn = Mock()
         mock_conn.uid.side_effect = Exception("Connection error")
 
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
+    def test_search_auth_error_reraises(self):
+        """Test re-raises auth errors so callers can handle reconnection."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = Exception("User not authenticated")
+
+        with pytest.raises(Exception, match="not authenticated"):
+            imap_common.get_message_ids_in_folder(mock_conn)
+
     def test_single_message(self):
         """Test fetching single message ID."""
         mock_conn = Mock()
@@ -707,8 +717,8 @@ class TestGetMessageIdsInFolder:
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
-    def test_fetch_exception_for_batch(self):
+    def test_fetch_non_auth_exception_continues(self):
-        """Test continues when fetch raises exception for a batch."""
+        """Test continues when fetch raises non-auth exception for a batch."""
         mock_conn = Mock()
         mock_conn.uid.side_effect = [
             ("OK", [b"1"]),
@@ -718,6 +728,17 @@ class TestGetMessageIdsInFolder:
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert result == {}
 
+    def test_fetch_auth_error_reraises(self):
+        """Test re-raises auth errors during fetch so callers can handle reconnection."""
+        mock_conn = Mock()
+        mock_conn.uid.side_effect = [
+            ("OK", [b"1"]),
+            Exception("AccessTokenExpired"),
+        ]
+
+        with pytest.raises(Exception, match="AccessTokenExpired"):
+            imap_common.get_message_ids_in_folder(mock_conn)
+
     def test_skips_empty_message_id(self):
         """Test that empty message IDs are not added to dict."""
         mock_conn = Mock()

test/test_imap_oauth2.py

@@ -28,9 +28,11 @@ def clear_oauth2_caches():
     """Clear module-level OAuth2 caches between tests."""
     imap_oauth2._msal_app_cache.clear()
     imap_oauth2._google_creds_cache.clear()
+    imap_oauth2._tenant_cache.clear()
     yield
     imap_oauth2._msal_app_cache.clear()
     imap_oauth2._google_creds_cache.clear()
+    imap_oauth2._tenant_cache.clear()
 
 
 class TestDetectOauth2Provider:
@@ -487,3 +489,101 @@ class TestRefreshOauth2Token:
         # The other two should see conf["oauth2_token"] changed and skip.
         assert call_count["value"] == 1
         assert conf["oauth2_token"] == "fresh_token"
+
+
+class TestIsTokenExpiredError:
+    """Tests for is_token_expired_error function."""
+
+    def test_detects_access_token_expired(self):
+        """Test detects 'accesstokenexpired' error."""
+        error = Exception("AUTHENTICATE failed: AccessTokenExpired")
+        assert imap_oauth2.is_token_expired_error(error) is True
+
+    def test_detects_session_invalidated(self):
+        """Test detects 'session invalidated' error."""
+        error = Exception("Session invalidated by server")
+        assert imap_oauth2.is_token_expired_error(error) is True
+
+    def test_case_insensitive_access_token_expired(self):
+        """Test case-insensitive matching for accesstokenexpired."""
+        error = Exception("ACCESSTOKENEXPIRED")
+        assert imap_oauth2.is_token_expired_error(error) is True
+
+    def test_case_insensitive_session_invalidated(self):
+        """Test case-insensitive matching for session invalidated."""
+        error = Exception("SESSION INVALIDATED")
+        assert imap_oauth2.is_token_expired_error(error) is True
+
+    def test_false_for_connection_error(self):
+        """Test returns False for connection errors."""
+        error = Exception("Connection refused")
+        assert imap_oauth2.is_token_expired_error(error) is False
+
+    def test_false_for_timeout_error(self):
+        """Test returns False for timeout errors."""
+        error = Exception("The read operation timed out")
+        assert imap_oauth2.is_token_expired_error(error) is False
+
+    def test_false_for_generic_error(self):
+        """Test returns False for generic errors."""
+        error = Exception("Something went wrong")
+        assert imap_oauth2.is_token_expired_error(error) is False
+
+    def test_false_for_not_authenticated(self):
+        """Test returns False for 'not authenticated' (handled by is_auth_error)."""
+        error = Exception("User not authenticated")
+        assert imap_oauth2.is_token_expired_error(error) is False
+
+
+class TestIsAuthError:
+    """Tests for is_auth_error function."""
+
+    def test_detects_access_token_expired(self):
+        """Test detects 'accesstokenexpired' error (via is_token_expired_error)."""
+        error = Exception("AUTHENTICATE failed: AccessTokenExpired")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_detects_session_invalidated(self):
+        """Test detects 'session invalidated' error (via is_token_expired_error)."""
+        error = Exception("Session invalidated by server")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_detects_not_authenticated(self):
+        """Test detects 'not authenticated' error."""
+        error = Exception("User not authenticated")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_detects_authentication_failed(self):
+        """Test detects 'authentication failed' error."""
+        error = Exception("Authentication failed for user@example.com")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_case_insensitive_not_authenticated(self):
+        """Test case-insensitive matching for not authenticated."""
+        error = Exception("NOT AUTHENTICATED")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_case_insensitive_authentication_failed(self):
+        """Test case-insensitive matching for authentication failed."""
+        error = Exception("AUTHENTICATION FAILED")
+        assert imap_oauth2.is_auth_error(error) is True
+
+    def test_false_for_connection_error(self):
+        """Test returns False for connection errors."""
+        error = Exception("Connection refused")
+        assert imap_oauth2.is_auth_error(error) is False
+
+    def test_false_for_timeout_error(self):
+        """Test returns False for timeout errors."""
+        error = Exception("The read operation timed out")
+        assert imap_oauth2.is_auth_error(error) is False
+
+    def test_false_for_generic_error(self):
+        """Test returns False for generic errors."""
+        error = Exception("Something went wrong")
+        assert imap_oauth2.is_auth_error(error) is False
+
+    def test_false_for_folder_not_found(self):
+        """Test returns False for folder errors."""
+        error = Exception("Folder not found: INBOX")
+        assert imap_oauth2.is_auth_error(error) is False

test/test_imap_session.py

@@ -0,0 +1,268 @@
+"""
+Tests for imap_session.py
+
+Tests cover:
+- ensure_connection() behavior with and without OAuth2
+- ensure_folder_session() connection change detection
+- Folder selection and re-selection logic
+- Error handling for connection and folder failures
+"""
+
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
+
+import imap_session
+
+
+class TestEnsureConnection:
+    """Tests for ensure_connection function."""
+
+    def test_with_oauth2_calls_refresh(self):
+        """Test that OAuth2 token refresh is called when oauth2 config is present."""
+        conf = {
+            "host": "imap.example.com",
+            "user": "user@example.com",
+            "password": "pass",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@example.com",
+            },
+        }
+        mock_conn = MagicMock()
+
+        with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
+            with patch.object(
+                imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn
+            ) as mock_ensure:
+                result = imap_session.ensure_connection(mock_conn, conf)
+
+        mock_refresh.assert_called_once_with(conf, "old_token")
+        mock_ensure.assert_called_once_with(mock_conn, conf)
+        assert result is mock_conn
+
+    def test_without_oauth2_skips_refresh(self):
+        """Test that token refresh is skipped when oauth2 config is not present."""
+        conf = {
+            "host": "imap.example.com",
+            "user": "user@example.com",
+            "password": "pass",
+            "oauth2": None,
+        }
+        mock_conn = MagicMock()
+
+        with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
+            with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
+                imap_session.ensure_connection(mock_conn, conf)
+
+        mock_refresh.assert_not_called()
+
+    def test_without_oauth2_key_skips_refresh(self):
+        """Test that token refresh is skipped when oauth2 key is missing entirely."""
+        conf = {
+            "host": "imap.example.com",
+            "user": "user@example.com",
+            "password": "pass",
+        }
+        mock_conn = MagicMock()
+
+        with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token") as mock_refresh:
+            with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
+                imap_session.ensure_connection(mock_conn, conf)
+
+        mock_refresh.assert_not_called()
+
+    def test_returns_connection_from_ensure_connection_from_conf(self):
+        """Test returns the connection from ensure_connection_from_conf."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+
+        with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
+            result = imap_session.ensure_connection(old_conn, conf)
+
+        assert result is new_conn
+
+    def test_returns_none_when_connection_fails(self):
+        """Test returns None when ensure_connection_from_conf fails."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+
+        with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=None):
+            result = imap_session.ensure_connection(MagicMock(), conf)
+
+        assert result is None
+
+    def test_with_none_connection(self):
+        """Test works correctly when passed None as connection."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        new_conn = MagicMock()
+
+        with patch.object(
+            imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn
+        ) as mock_ensure:
+            result = imap_session.ensure_connection(None, conf)
+
+        mock_ensure.assert_called_once_with(None, conf)
+        assert result is new_conn
+
+
+class TestEnsureFolderSession:
+    """Tests for ensure_folder_session function."""
+
+    def test_same_connection_no_folder_select(self):
+        """Test that folder is not re-selected when connection stays the same."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        mock_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=mock_conn):
+            result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is mock_conn
+        assert success is True
+        mock_conn.select.assert_not_called()
+
+    def test_new_connection_selects_folder(self):
+        """Test that folder is selected when connection changes."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is new_conn
+        assert success is True
+        new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
+
+    def test_new_connection_selects_folder_readwrite(self):
+        """Test that folder is selected with readonly=False when specified."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(
+                old_conn, conf, "[Gmail]/All Mail", readonly=False
+            )
+
+        assert result_conn is new_conn
+        assert success is True
+        new_conn.select.assert_called_once_with('"[Gmail]/All Mail"', readonly=False)
+
+    def test_none_connection_returns_failure(self):
+        """Test returns (None, False) when ensure_connection returns None."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        mock_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=None):
+            result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is None
+        assert success is False
+
+    def test_folder_select_failure_returns_false(self):
+        """Test returns (conn, False) when folder selection fails."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+        new_conn.select.side_effect = Exception("Folder not found")
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "NonExistent", readonly=True)
+
+        assert result_conn is new_conn
+        assert success is False
+
+    def test_none_old_connection_selects_folder(self):
+        """Test that folder is selected when old connection was None."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        new_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(None, conf, "INBOX", readonly=True)
+
+        assert result_conn is new_conn
+        assert success is True
+        new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
+
+    def test_folder_name_with_spaces(self):
+        """Test folder selection with folder names containing spaces."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "My Folder", readonly=True)
+
+        assert success is True
+        new_conn.select.assert_called_once_with('"My Folder"', readonly=True)
+
+    def test_folder_select_imap_error(self):
+        """Test handles IMAP-specific errors during folder selection."""
+        conf = {"host": "imap.example.com", "user": "user", "password": "pass"}
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+        new_conn.select.side_effect = Exception("IMAP error: NO SELECT failed")
+
+        with patch.object(imap_session, "ensure_connection", return_value=new_conn):
+            result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is new_conn
+        assert success is False
+
+
+class TestEnsureFolderSessionWithOAuth2:
+    """Integration-style tests for ensure_folder_session with OAuth2."""
+
+    def test_oauth2_refresh_triggers_folder_reselect(self):
+        """Test that OAuth2 token refresh (new connection) triggers folder reselection."""
+        conf = {
+            "host": "outlook.office365.com",
+            "user": "user@example.com",
+            "password": "oauth2_token_here",
+            "oauth2_token": "old_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@example.com",
+            },
+        }
+        old_conn = MagicMock()
+        new_conn = MagicMock()
+
+        # Simulate token refresh causing a new connection
+        with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
+            with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=new_conn):
+                result_conn, success = imap_session.ensure_folder_session(old_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is new_conn
+        assert success is True
+        new_conn.select.assert_called_once_with('"INBOX"', readonly=True)
+
+    def test_healthy_connection_no_reselect(self):
+        """Test that healthy connection (same object) doesn't reselect folder."""
+        conf = {
+            "host": "outlook.office365.com",
+            "user": "user@example.com",
+            "password": "oauth2_token_here",
+            "oauth2_token": "valid_token",
+            "oauth2": {
+                "provider": "microsoft",
+                "client_id": "client-id",
+                "email": "user@example.com",
+            },
+        }
+        mock_conn = MagicMock()
+
+        # Simulate healthy connection (same object returned)
+        with patch.object(imap_session.imap_oauth2, "refresh_oauth2_token"):
+            with patch.object(imap_session.imap_common, "ensure_connection_from_conf", return_value=mock_conn):
+                result_conn, success = imap_session.ensure_folder_session(mock_conn, conf, "INBOX", readonly=True)
+
+        assert result_conn is mock_conn
+        assert success is True
+        mock_conn.select.assert_not_called()