feat(tasks): refactor IMAP and watch folder polling to support multi-tenant user attribution

- Add owner_id parameter to pull_inbox() and fetch_attachments_and_enqueue()
  to attribute ingested documents to the correct user
- Add _pull_user_integration_imap() to poll IMAP sources from UserIntegration model
- Add _pull_user_integration_watch_folders() to scan watch folders from UserIntegration model
- Add _is_safe_watch_path() for path traversal security on user-configured paths
- Add _scan_user_watch_folder() that passes owner_id to _enqueue_file()
- Update _enqueue_file() to forward owner_id to process_document/convert_to_pdf
- Update celery beat schedule to always enable IMAP and watch folder polling
  (user integrations can exist without system-level config)
- Ensure individual connection failures don't crash the polling loop
- Update existing tests for new function signatures

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>

app/celery_worker.py

@@ -63,15 +63,13 @@ def test_task():
 check_credentials.apply_async(countdown=10)  # Run 10 seconds after worker starts
 
 celery.conf.beat_schedule = {
-    "poll-inboxes-every-minute": (
+    # IMAP polling — always enabled because per-user IMAP integrations may
-        {
+    # exist in the database even when no system-level IMAP hosts are configured.
-            "task": "app.tasks.imap_tasks.pull_all_inboxes",
+    "poll-inboxes-every-minute": {
-            "schedule": crontab(minute="*/1"),  # every 1 minute
+        "task": "app.tasks.imap_tasks.pull_all_inboxes",
-            "options": {"expires": 55},  # Ensure tasks don't pile up
+        "schedule": crontab(minute="*/1"),  # every 1 minute
-        }
+        "options": {"expires": 55},  # Ensure tasks don't pile up
-        if (settings.imap1_host or settings.imap2_host)
+    },
-        else None
-    ),
     # Add Uptime Kuma ping task if configured
     "ping-uptime-kuma": (
         {
@@ -100,27 +98,15 @@ celery.conf.beat_schedule = {
         "schedule": crontab(minute="*/1"),  # Every minute
         "options": {"expires": 55},  # Must complete within 55 seconds
     },
-    # Watch folder scanning — polls local paths, FTP, SFTP, and cloud ingest folders.
+    # Watch folder scanning — always enabled because per-user WATCH_FOLDER
+    # integrations may exist in the database even when no system-level watch
+    # folder settings are configured.
     # Schedule is controlled by WATCH_FOLDER_POLL_INTERVAL (default: 1 minute).
-    "scan-watch-folders": (
+    "scan-watch-folders": {
-        {
+        "task": "app.tasks.watch_folder_tasks.scan_all_watch_folders",
-            "task": "app.tasks.watch_folder_tasks.scan_all_watch_folders",
+        "schedule": crontab(minute=f"*/{max(1, settings.watch_folder_poll_interval)}"),
-            "schedule": crontab(minute=f"*/{max(1, settings.watch_folder_poll_interval)}"),
+        "options": {"expires": 55},
-            "options": {"expires": 55},
+    },
-        }
-        if (
-            settings.watch_folders
-            or settings.ftp_ingest_enabled
-            or settings.sftp_ingest_enabled
-            or settings.dropbox_ingest_enabled
-            or settings.google_drive_ingest_enabled
-            or settings.onedrive_ingest_enabled
-            or settings.nextcloud_ingest_enabled
-            or settings.s3_ingest_enabled
-            or settings.webdav_ingest_enabled
-        )
-        else None
-    ),
     # Backfill embeddings for files that were processed before the
     # embedding pipeline was enabled, or where the embedding task failed.
     "backfill-missing-embeddings": {

app/tasks/imap_tasks.py

@@ -149,6 +149,9 @@ def pull_all_inboxes():
         # Per-user IMAP accounts from the database
         _pull_user_imap_accounts()
 
+        # Per-user IMAP integrations from the UserIntegration model
+        _pull_user_integration_imap()
+
         logger.info("Finished pull_all_inboxes")
 
     finally:
@@ -176,6 +179,7 @@ def _pull_user_imap_accounts() -> None:
                         password=_decrypt_imap_password(acct.password),
                         use_ssl=acct.use_ssl,
                         delete_after_process=acct.delete_after_process,
+                        owner_id=acct.owner_id,
                     )
                     # Record successful poll
                     acct.last_checked_at = datetime.now(timezone.utc)
@@ -202,6 +206,91 @@ def _pull_user_imap_accounts() -> None:
         logger.error("Failed to process per-user IMAP accounts: %s", exc)
 
 
+def _pull_user_integration_imap() -> None:
+    """Iterate over all active IMAP UserIntegrations and pull their inboxes.
+
+    This polls the ``user_integrations`` table for records with
+    ``integration_type='IMAP'``, ``direction='SOURCE'``, and ``is_active=True``.
+    Each integration's config/credentials are decoded and passed to
+    :func:`pull_inbox` with the owning user's ``owner_id`` so that ingested
+    documents are correctly attributed.
+
+    Individual connection failures are caught and recorded on the integration
+    without crashing the polling loop.
+    """
+    try:
+        import json as _json
+
+        from app.models import IntegrationDirection, IntegrationType, UserIntegration
+        from app.utils.encryption import decrypt_value
+
+        db = _get_db_session()
+        try:
+            integrations = (
+                db.query(UserIntegration)
+                .filter(
+                    UserIntegration.integration_type == IntegrationType.IMAP,
+                    UserIntegration.direction == IntegrationDirection.SOURCE,
+                    UserIntegration.is_active.is_(True),
+                )
+                .all()
+            )
+            logger.info("Processing %d IMAP UserIntegration(s)", len(integrations))
+            for integ in integrations:
+                account_identifier = f"integration_{integ.owner_id}_{integ.id}"
+                try:
+                    cfg = _json.loads(integ.config) if integ.config else {}
+                    raw_creds = decrypt_value(integ.credentials) if integ.credentials else None
+                    creds = _json.loads(raw_creds) if raw_creds else {}
+
+                    host = cfg.get("host")
+                    port = int(cfg.get("port", 993))
+                    username = cfg.get("username")
+                    password = creds.get("password")
+                    use_ssl = cfg.get("use_ssl", True)
+                    delete_after = cfg.get("delete_after_process", False)
+
+                    if not (host and username and password):
+                        logger.warning(
+                            "IMAP integration %d (owner %s) has incomplete config — skipping.",
+                            integ.id,
+                            integ.owner_id,
+                        )
+                        continue
+
+                    pull_inbox(
+                        mailbox_key=account_identifier,
+                        host=host,
+                        port=port,
+                        username=username,
+                        password=password,
+                        use_ssl=use_ssl,
+                        delete_after_process=delete_after,
+                        owner_id=integ.owner_id,
+                    )
+                    integ.last_used_at = datetime.now(timezone.utc)
+                    integ.last_error = None
+                    db.commit()
+                except Exception as exc:  # noqa: BLE001
+                    error_msg = str(exc)[:_MAX_ERROR_LENGTH]
+                    logger.error(
+                        "Error pulling IMAP integration %d (owner %s): %s",
+                        integ.id,
+                        integ.owner_id,
+                        error_msg,
+                    )
+                    try:
+                        integ.last_used_at = datetime.now(timezone.utc)
+                        integ.last_error = error_msg
+                        db.commit()
+                    except Exception:  # noqa: BLE001
+                        db.rollback()
+        finally:
+            db.close()
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Failed to process IMAP UserIntegrations: %s", exc)
+
+
 def check_and_pull_mailbox(
     mailbox_key: str,
     host: str | None,
@@ -228,7 +317,7 @@ def check_and_pull_mailbox(
     )
 
 
-def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_after_process):
+def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_after_process, owner_id=None):
     """
     Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
     and processes attachments while preserving the original unread status.
@@ -238,6 +327,10 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
       - Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment".
 
     For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.
+
+    Args:
+        owner_id: Optional user identifier. When provided, ingested documents are
+                  attributed to this user via ``process_document`` / ``convert_to_pdf``.
     """
     logger.info("Connecting to %s at %s:%s (SSL=%s)", mailbox_key, host, port, use_ssl)
     processed_emails = load_processed_emails()
@@ -300,7 +393,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
 
             # Process attachments (and convert non-PDF files).
             # We call the function without assigning its return value since it is not used.
-            fetch_attachments_and_enqueue(email_message)
+            fetch_attachments_and_enqueue(email_message, owner_id=owner_id)
 
             if settings.imap_readonly_mode:
                 logger.info("Readonly mode: skipping mailbox modifications for %s in %s", msg_id, mailbox_key)
@@ -329,7 +422,7 @@ def pull_inbox(mailbox_key, host, port, username, password, use_ssl, delete_afte
         logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)
 
 
-def fetch_attachments_and_enqueue(email_message):
+def fetch_attachments_and_enqueue(email_message, owner_id: str | None = None):
     """
     Extracts attachments from the email and processes only allowed file types.
 
@@ -354,6 +447,11 @@ def fetch_attachments_and_enqueue(email_message):
     If the attachment is a PDF (by extension or MIME type), it is enqueued for upload;
     any other allowed file is enqueued for conversion to PDF.
 
+    Args:
+        email_message: The parsed email message to extract attachments from.
+        owner_id: Optional user identifier forwarded to ``process_document`` /
+                  ``convert_to_pdf`` for multi-tenant attribution.
+
     Returns True if at least one allowed attachment was processed.
     """
     has_attachment = False
@@ -381,11 +479,11 @@ def fetch_attachments_and_enqueue(email_message):
 
         # If it's a PDF by MIME type or extension, process it directly
         if mime_type == "application/pdf" or is_pdf_by_extension:
-            process_document.delay(file_path)
+            process_document.delay(file_path, owner_id=owner_id)
             logger.info("Enqueued PDF for upload: %s (MIME: %s)", filename, mime_type)
         elif mime_type in ALLOWED_MIME_TYPES:
             # Other allowed files are sent for conversion
-            convert_to_pdf.delay(file_path)
+            convert_to_pdf.delay(file_path, owner_id=owner_id)
             logger.info("Enqueued file for conversion to PDF: %s", filename)
 
         has_attachment = True

app/tasks/watch_folder_tasks.py

@@ -39,8 +39,27 @@ FTP_INGEST_CACHE_FILE = os.path.join(settings.workdir, "ftp_ingest_processed.jso
 # Cache file for tracking already-ingested files (SFTP watch folder)
 SFTP_INGEST_CACHE_FILE = os.path.join(settings.workdir, "sftp_ingest_processed.json")
 
+# Per-integration watch-folder cache file prefix
+_USER_WF_CACHE_PREFIX = os.path.join(settings.workdir, "user_wf_")
+
 _CACHE_RETENTION_DAYS = 30
 
+# Maximum length to store as last_error on UserIntegration to prevent DB bloat
+_MAX_ERROR_LENGTH = 500
+
+# Database session factory (imported lazily to avoid circular imports)
+_db_session_factory = None
+
+
+def _get_db_session():
+    """Return a new SQLAlchemy session (lazy import to avoid startup issues)."""
+    global _db_session_factory  # noqa: PLW0603
+    if _db_session_factory is None:
+        from app.database import SessionLocal
+
+        _db_session_factory = SessionLocal
+    return _db_session_factory()
+
 
 # ---------------------------------------------------------------------------
 # Locking helpers
@@ -122,17 +141,24 @@ def _is_allowed_file(filename: str) -> bool:
     return ext in ALLOWED_EXTENSIONS or filename.lower().endswith(".pdf")
 
 
-def _enqueue_file(file_path: str, *, filename: str | None = None) -> None:
+def _enqueue_file(file_path: str, *, filename: str | None = None, owner_id: str | None = None) -> None:
-    """Enqueue a local file path for document processing."""
+    """Enqueue a local file path for document processing.
+
+    Args:
+        file_path: Absolute path to the file on disk.
+        filename: Optional display filename (defaults to basename of *file_path*).
+        owner_id: Optional user identifier forwarded to ``process_document`` /
+                  ``convert_to_pdf`` for multi-tenant attribution.
+    """
     fname = filename or os.path.basename(file_path)
     _, ext = os.path.splitext(fname)
     mime_check = ext.lower() in {".pdf"}
 
     if mime_check or fname.lower().endswith(".pdf"):
-        process_document.delay(file_path)
+        process_document.delay(file_path, owner_id=owner_id)
         logger.info("Enqueued for processing: %s", fname)
     else:
-        convert_to_pdf.delay(file_path)
+        convert_to_pdf.delay(file_path, owner_id=owner_id)
         logger.info("Enqueued for PDF conversion: %s", fname)
 
 
@@ -1288,6 +1314,216 @@ def scan_webdav_watch_folder() -> dict:
     return {"status": "ok", "files_enqueued": n, "folder": ingest_folder}
 
 
+# ---------------------------------------------------------------------------
+# Per-user watch folder integration scanning
+# ---------------------------------------------------------------------------
+
+
+def _is_safe_watch_path(folder_path: str) -> bool:
+    """Validate that a user-configured watch folder path is safe.
+
+    Rejects paths that attempt directory traversal (``..``), use relative
+    references, or are not absolute.  This prevents a malicious user from
+    configuring a watch folder that could escape its intended directory.
+
+    Args:
+        folder_path: The path to validate.
+
+    Returns:
+        ``True`` if the path is considered safe, ``False`` otherwise.
+    """
+    if not folder_path:
+        return False
+    # Must be absolute
+    if not os.path.isabs(folder_path):
+        logger.warning("Rejecting non-absolute watch folder path: %s", folder_path)
+        return False
+    # Resolve to canonical path and ensure no traversal components exist
+    resolved = os.path.realpath(folder_path)
+    if ".." in folder_path.split(os.sep):
+        logger.warning("Rejecting path with traversal components: %s", folder_path)
+        return False
+    # Ensure resolved path matches the original intent (no symlink escapes)
+    if resolved != os.path.normpath(folder_path):
+        logger.warning(
+            "Watch folder path resolves differently (possible symlink escape): %s -> %s",
+            folder_path,
+            resolved,
+        )
+        return False
+    return True
+
+
+def _scan_user_watch_folder(
+    folder_path: str,
+    cache: dict[str, str],
+    delete_after: bool,
+    owner_id: str,
+) -> int:
+    """Scan a user-configured local watch folder, attributing files to *owner_id*.
+
+    Delegates to the same file-scanning logic as :func:`_scan_local_folder` but
+    passes ``owner_id`` to :func:`_enqueue_file` so that ingested documents are
+    correctly attributed to the user.
+
+    Args:
+        folder_path: Absolute directory path to scan.
+        cache: In-memory dict of already-processed file keys.
+        delete_after: Whether to remove the source file after ingestion.
+        owner_id: The user to attribute ingested documents to.
+
+    Returns:
+        Number of files newly enqueued.
+    """
+    if not os.path.isdir(folder_path):
+        logger.warning("User watch folder does not exist or is not a directory: %s", folder_path)
+        return 0
+
+    count = 0
+    try:
+        entries = os.scandir(folder_path)
+    except PermissionError as exc:
+        logger.error("Cannot scan user watch folder %s: %s", folder_path, exc)
+        return 0
+
+    for entry in entries:
+        if not entry.is_file(follow_symlinks=True):
+            continue
+        if not _is_allowed_file(entry.name):
+            continue
+
+        abs_path = entry.path
+        if abs_path in cache:
+            continue
+
+        dest_filename = f"uwf_{owner_id}_{entry.name}"
+        dest_path = os.path.join(settings.workdir, dest_filename)
+        if os.path.exists(dest_path):
+            base, ext2 = os.path.splitext(dest_filename)
+            dest_path = os.path.join(settings.workdir, f"{base}_{int(datetime.now().timestamp())}{ext2}")
+
+        try:
+            import shutil
+
+            shutil.copy2(abs_path, dest_path)
+        except OSError as exc:
+            logger.error("Failed to copy %s to workdir: %s", abs_path, exc)
+            continue
+
+        _enqueue_file(dest_path, owner_id=owner_id)
+        _mark_processed(cache, abs_path)
+        count += 1
+
+        if delete_after:
+            try:
+                os.remove(abs_path)
+                logger.info("Deleted source file after ingestion: %s", abs_path)
+            except OSError as exc:
+                logger.warning("Could not delete source file %s: %s", abs_path, exc)
+
+    return count
+
+
+def _pull_user_integration_watch_folders() -> dict:
+    """Iterate over all active WATCH_FOLDER UserIntegrations and scan their paths.
+
+    Polls the ``user_integrations`` table for records with
+    ``integration_type='WATCH_FOLDER'``, ``direction='SOURCE'``, and
+    ``is_active=True``.  Each integration's config is decoded and the
+    configured ``folder_path`` is scanned for new files, which are enqueued
+    with the owning user's ``owner_id``.
+
+    Path traversal protection is enforced on the configured path.
+
+    Individual integration failures are caught and recorded without crashing
+    the polling loop.
+
+    Returns:
+        Summary dict with ``status`` and ``integrations_processed`` count.
+    """
+    try:
+        import json as _json
+
+        from app.models import IntegrationDirection, IntegrationType, UserIntegration
+
+        db = _get_db_session()
+        try:
+            integrations = (
+                db.query(UserIntegration)
+                .filter(
+                    UserIntegration.integration_type == IntegrationType.WATCH_FOLDER,
+                    UserIntegration.direction == IntegrationDirection.SOURCE,
+                    UserIntegration.is_active.is_(True),
+                )
+                .all()
+            )
+            logger.info("Processing %d WATCH_FOLDER UserIntegration(s)", len(integrations))
+            total_files = 0
+            for integ in integrations:
+                try:
+                    cfg = _json.loads(integ.config) if integ.config else {}
+                    folder_path = cfg.get("folder_path", "")
+                    delete_after = cfg.get("delete_after_process", False)
+
+                    if not folder_path:
+                        logger.warning(
+                            "Watch folder integration %d (owner %s) has no folder_path — skipping.",
+                            integ.id,
+                            integ.owner_id,
+                        )
+                        continue
+
+                    if not _is_safe_watch_path(folder_path):
+                        error_msg = f"Unsafe watch folder path rejected: {folder_path}"
+                        logger.error(
+                            "Watch folder integration %d (owner %s): %s",
+                            integ.id,
+                            integ.owner_id,
+                            error_msg,
+                        )
+                        integ.last_error = error_msg[:_MAX_ERROR_LENGTH]
+                        db.commit()
+                        continue
+
+                    cache_file = f"{_USER_WF_CACHE_PREFIX}{integ.id}.json"
+                    cache = _load_cache(cache_file)
+                    n = _scan_user_watch_folder(folder_path, cache, delete_after, integ.owner_id)
+                    _save_cache(cache_file, cache)
+
+                    total_files += n
+                    integ.last_used_at = datetime.now(timezone.utc)
+                    integ.last_error = None
+                    db.commit()
+                    logger.info(
+                        "Watch folder integration %d (owner %s): %d file(s) enqueued from %s",
+                        integ.id,
+                        integ.owner_id,
+                        n,
+                        folder_path,
+                    )
+                except Exception as exc:  # noqa: BLE001
+                    error_msg = str(exc)[:_MAX_ERROR_LENGTH]
+                    logger.error(
+                        "Error scanning watch folder integration %d (owner %s): %s",
+                        integ.id,
+                        integ.owner_id,
+                        error_msg,
+                    )
+                    try:
+                        integ.last_used_at = datetime.now(timezone.utc)
+                        integ.last_error = error_msg
+                        db.commit()
+                    except Exception:  # noqa: BLE001
+                        db.rollback()
+
+            return {"status": "ok", "integrations_processed": len(integrations), "files_enqueued": total_files}
+        finally:
+            db.close()
+    except Exception as exc:  # noqa: BLE001
+        logger.error("Failed to process WATCH_FOLDER UserIntegrations: %s", exc)
+        return {"status": "error", "reason": str(exc)[:_MAX_ERROR_LENGTH]}
+
+
 @shared_task
 def scan_all_watch_folders() -> dict:
     """
@@ -1303,6 +1539,7 @@ def scan_all_watch_folders() -> dict:
     7. Nextcloud ingest folder (if enabled)
     8. Amazon S3 ingest prefix (if enabled)
     9. WebDAV ingest folder (if enabled)
+    10. Per-user WATCH_FOLDER integrations from the database
     """
     if not _acquire_lock(WATCH_FOLDER_LOCK_KEY):
         logger.info("Watch folder scan already running — skipping this cycle.")
@@ -1319,6 +1556,7 @@ def scan_all_watch_folders() -> dict:
         results["nextcloud"] = scan_nextcloud_watch_folder()
         results["s3"] = scan_s3_watch_folder()
         results["webdav"] = scan_webdav_watch_folder()
+        results["user_watch_folders"] = _pull_user_integration_watch_folders()
     finally:
         _release_lock(WATCH_FOLDER_LOCK_KEY)
 

tests/test_celery_worker.py

@@ -163,10 +163,10 @@ class TestBeatScheduleConfiguration:
 
     # -- Conditional schedule entries ----------------------------------------
 
-    def test_imap_schedule_absent_when_not_configured(self):
+    def test_imap_schedule_always_present(self):
-        """Test IMAP polling absent when neither imap host is set."""
+        """Test IMAP polling is always scheduled (user integrations may exist)."""
         mod = _reload_celery_worker()
-        assert "poll-inboxes-every-minute" not in mod.celery.conf.beat_schedule
+        assert "poll-inboxes-every-minute" in mod.celery.conf.beat_schedule
 
     def test_uptime_kuma_schedule_absent_when_not_configured(self):
         """Test Uptime Kuma ping absent when url is not set."""

tests/test_watch_folder_tasks.py

@@ -834,7 +834,7 @@ class TestEnqueueFileNonPdf:
             patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
         ):
             _enqueue_file("/tmp/doc.docx")
-        mock_conv.delay.assert_called_once_with("/tmp/doc.docx")
+        mock_conv.delay.assert_called_once_with("/tmp/doc.docx", owner_id=None)
         mock_proc.delay.assert_not_called()
 
     def test_pdf_triggers_process_document(self):
@@ -846,7 +846,7 @@ class TestEnqueueFileNonPdf:
             patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
         ):
             _enqueue_file("/tmp/report.pdf")
-        mock_proc.delay.assert_called_once_with("/tmp/report.pdf")
+        mock_proc.delay.assert_called_once_with("/tmp/report.pdf", owner_id=None)
         mock_conv.delay.assert_not_called()
 
     def test_pdf_detected_via_filename_param(self):
@@ -858,7 +858,7 @@ class TestEnqueueFileNonPdf:
             patch("app.tasks.watch_folder_tasks.process_document") as mock_proc,
         ):
             _enqueue_file("/tmp/somefile", filename="renamed.pdf")
-        mock_proc.delay.assert_called_once_with("/tmp/somefile")
+        mock_proc.delay.assert_called_once_with("/tmp/somefile", owner_id=None)
         mock_conv.delay.assert_not_called()