deer-flow/backend/app/gateway/upload_ingestion.py

"""Shared thread-upload ingestion service (Projects Phase 2 Slice C, §7.3 item 3).

Owns the whole pipeline for files landing in a thread's uploads directory:
staging, ``claim_unique_filename``, size checks, optional conversion under
``uploads.auto_convert_documents``, sandbox-readable permissions, and
non-mounted-provider synchronization of original + derived files through the
authorized sandbox request lease. A denied ``sandbox:execute`` retains the
host upload without allocating a sandbox (ordinary uploads behavior);
acquisition, sync and conversion failures follow the ordinary upload error
and cleanup behavior.

Two callers share this pipeline:

- the ordinary uploads endpoint (``routers/uploads.py``), a thin adapter
  mapping ``UploadFile`` parts to chunk streams with zero behavior change;
- the project-shelf attach route (``routers/project_documents.py``), which
  feeds a staged shelf copy through the same lifecycle (§7.3 item 3).

Behavior-preservation bridge: the pre-extraction upload tests pin this
pipeline by patching collaborators on the ``routers.uploads`` module
namespace (``ensure_uploads_dir``, ``get_sandbox_provider``,
``_get_upload_limits``, ``_auto_convert_documents_enabled``,
``convert_file_to_markdown``, the low-level file machinery). The service
therefore resolves every collaborator through that module object at call
time (:func:`_uploads`) instead of importing the names directly, so those
patches keep binding to the one pipeline both callers use.
"""

from __future__ import annotations

import logging
import uuid
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any

from fastapi import HTTPException

from deerflow.config.app_config import AppConfig
from deerflow.utils.file_io import run_file_io

if TYPE_CHECKING:
    from fastapi import Request

    from app.gateway.authz import SandboxRequestLease
    from app.gateway.routers.uploads import UploadLimits

logger = logging.getLogger(__name__)


def _uploads() -> Any:
    """Return the uploads router module (late binding — see module docstring)."""
    from app.gateway.routers import uploads

    return uploads


class UnsafeFilenameError(ValueError):
    """The display filename could not be normalized/claimed; ordinary behavior skips it."""


class UnsafeUploadDestinationError(Exception):
    """The claimed destination failed safety validation; carries the claimed name.

    Ordinary uploads record the file as skipped (``skipped_files``) and keep
    ingesting the rest of the request; attach maps it to an ingestion
    failure. Wraps the manager's ``UnsafeUploadPathError`` so the claimed
    (deduplicated) filename survives to the caller.
    """

    def __init__(self, filename: str) -> None:
        super().__init__(f"Unsafe upload destination: {filename}")
        self.filename = filename


class ThreadUploadIngestionService:
    """One thread's ingestion session: staging → conversion → permissions → sync.

    Lifecycle: :meth:`open` (uploads dir + existing-name seed + optional
    sandbox lease), one or
    more :meth:`ingest_chunks` calls, :meth:`finalize` (sandbox-readable
    permissions + non-mounted provider sync), then :meth:`aclose` (lease
    release) — use it as an async context manager. :meth:`cleanup_written`
    removes every file written by this session; the ordinary endpoint invokes
    it when a per-file failure aborts the request, and attach mirrors that.
    """

    def __init__(self, *, request: Request | None, thread_id: str, user_id: str, app_config: AppConfig) -> None:
        self._request = request
        self._thread_id = thread_id
        self._user_id = user_id
        self._config = app_config
        self._limits: UploadLimits | None = None
        self._uploads_dir: Path | None = None
        self._sync_to_sandbox = False
        self._sandbox_lease: SandboxRequestLease | None = None
        self._sandbox: Any = None
        self._auto_convert = False
        self._seen_filenames: set[str] = set()
        self._written_paths: list[Path] = []
        self._sync_targets: list[tuple[Path, str]] = []
        self._total_size = 0

    async def __aenter__(self) -> ThreadUploadIngestionService:
        await self.open()
        return self

    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
        await self.aclose()

    @property
    def limits(self) -> UploadLimits:
        assert self._limits is not None, "open() must run before limits is read"
        return self._limits

    async def open(self) -> None:
        """Resolve limits, ensure the uploads dir, seed existing names, acquire the lease.

        Mirrors the ordinary endpoint's setup order: a role denied
        ``sandbox:execute`` skips acquisition entirely (the host upload still
        succeeds); an allowed caller whose sandbox vanishes right after
        acquiring is a 500. ``ValueError`` from the uploads-dir resolution
        (unsafe thread id) propagates for the caller to map. The existing-name
        seed makes every later claim unique against the thread's CURRENT
        files, so no ingestion — ordinary upload or shelf attach — ever
        silently replaces a conversation file (same-name re-uploads land as
        ``name_1.ext``); the directory scan is offloaded per the filesystem
        convention.
        """
        uploads = _uploads()
        self._limits = uploads._get_upload_limits(self._config)
        self._uploads_dir = await run_file_io(uploads.ensure_uploads_dir, self._thread_id, user_id=self._user_id)
        listing = await run_file_io(uploads.list_files_in_dir, self._uploads_dir)
        self._seen_filenames.update(entry["filename"] for entry in listing["files"])
        sandbox_provider = uploads.get_sandbox_provider()
        self._sync_to_sandbox = not uploads._uses_thread_data_mounts(sandbox_provider)
        if self._sync_to_sandbox:
            self._sandbox_lease = await uploads.try_acquire_sandbox_for_request(
                self._request,
                sandbox_provider,
                self._thread_id,
                user_id=self._user_id,
                app_config=self._config,
                owner_prefix="gateway:upload",
                release_on_last=False,
            )
            self._sandbox = self._sandbox_lease.sandbox
            if not self._sandbox_lease.denied and self._sandbox is None:
                raise HTTPException(status_code=500, detail="Failed to acquire sandbox")
        self._auto_convert = uploads._auto_convert_documents_enabled(self._config)

    async def _link_commit_with_retry(self, uploads: Any, staged_path: Path, claimed_name: str) -> tuple[str, Path]:
        """Publish a staged file under *claimed_name*, atomically and with no overwrite.

        ``os.link`` (inside ``_commit_upload_temp_no_overwrite`` /
        ``_link_staged_no_overwrite``) fails with :class:`FileExistsError`
        when a concurrent session won the name after this session's seed —
        the next suffix is claimed against the seeded set and the link
        retried, so two concurrent ingestions of one name always land as
        ``name.ext`` + ``name_1.ext`` with both byte streams intact. The
        staged bytes stay hidden under the ``.upload-*.part`` pattern until
        the link makes the final name visible, fully formed.
        """
        assert self._uploads_dir is not None, "open() must run before committing destinations"
        name = claimed_name
        while True:
            try:
                return name, await run_file_io(uploads._link_staged_no_overwrite, staged_path, self._uploads_dir, name)
            except FileExistsError:
                name = uploads.claim_unique_filename(name, self._seen_filenames)

    async def ingest_chunks(self, chunks: AsyncIterator[bytes], *, display_name: str) -> dict[str, Any]:
        """Ingest one file from a chunk stream; return its wire metadata dict.

        Owns the per-file pipeline exactly as the ordinary endpoint defines
        it: normalize + claim a unique name, staged write with single-file
        and request-total size caps (413), an atomic no-overwrite link commit
        with next-suffix retry, and conversion with its companion-name claim
        when ``uploads.auto_convert_documents`` is on. Raises
        :class:`UnsafeFilenameError` (caller logs and skips),
        :class:`UnsafeUploadDestinationError` (caller records a skipped
        file), ``HTTPException`` (size/over-limit — caller cleans up and
        re-raises) or an unexpected error (caller cleans up and answers 500).
        """
        assert self._uploads_dir is not None and self._limits is not None, "open() must run before ingest_chunks()"
        uploads = _uploads()
        try:
            original_filename = uploads.normalize_filename(display_name)
        except ValueError as exc:
            raise UnsafeFilenameError(str(exc)) from exc
        safe_filename = uploads.claim_unique_filename(original_filename, self._seen_filenames)

        file_size = 0
        upload_temp = None
        try:
            upload_temp = await run_file_io(uploads._prepare_upload_destination, self._uploads_dir, safe_filename)
            async for chunk in chunks:
                file_size += len(chunk)
                self._total_size += len(chunk)
                if file_size > self._limits.max_file_size:
                    raise HTTPException(status_code=413, detail=f"File too large: {safe_filename}")
                if self._total_size > self._limits.max_total_size:
                    raise HTTPException(status_code=413, detail="Total upload size too large")
                await run_file_io(uploads._write_upload_chunk, upload_temp, chunk)
            # Link-commit with collision retry: the FileExistsError arm
            # leaves the staged part in place for the retry under the next
            # suffix (the handle's second close is idempotent).
            while True:
                try:
                    file_path = await run_file_io(uploads._commit_upload_temp_no_overwrite, upload_temp, self._uploads_dir, safe_filename)
                    break
                except FileExistsError:
                    safe_filename = uploads.claim_unique_filename(safe_filename, self._seen_filenames)
            upload_temp = None
        except uploads.UnsafeUploadPathError as exc:
            if upload_temp is not None:
                await run_file_io(uploads._abort_upload_temp, upload_temp)
            raise UnsafeUploadDestinationError(safe_filename) from exc
        except Exception:
            if upload_temp is not None:
                await run_file_io(uploads._abort_upload_temp, upload_temp)
            raise

        self._written_paths.append(file_path)
        virtual_path = uploads.upload_virtual_path(safe_filename)
        if self._sync_to_sandbox:
            self._sync_targets.append((file_path, virtual_path))

        file_info: dict[str, Any] = {
            "filename": safe_filename,
            "size": file_size,
            "path": str(self._uploads_dir / safe_filename),
            "virtual_path": virtual_path,
            "artifact_url": uploads.upload_artifact_url(self._thread_id, safe_filename),
        }
        if safe_filename != original_filename:
            file_info["original_filename"] = original_filename
        logger.info(f"Saved file: {safe_filename} ({file_size} bytes) to {file_info['path']}")

        if self._auto_convert and file_path.suffix.lower() in uploads.CONVERTIBLE_EXTENSIONS:
            # The companion gets the same atomic no-overwrite commit as the
            # original: staged under the hidden .part pattern, link-committed
            # with next-suffix retry — conversion can never silently truncate
            # another uploaded or derived file, in this session or a
            # concurrent one.
            provisional_md_name = Path(safe_filename).with_suffix(".md").name
            unique_md_name = uploads.claim_unique_filename(provisional_md_name, self._seen_filenames)
            md_staging = self._uploads_dir / f"{uploads.UPLOAD_STAGING_PREFIX}{uuid.uuid4().hex}{uploads.UPLOAD_STAGING_SUFFIX}"
            try:
                md_staged = await uploads.convert_file_to_markdown(file_path, output_path=md_staging)
            except Exception:
                self._seen_filenames.discard(unique_md_name)
                await run_file_io(md_staging.unlink, True)
                raise
            if not md_staged:
                # Conversion failed and wrote nothing (or a partial staged
                # file, removed here): release the claim; holding it would
                # rename a later same-stem upload against a name nothing
                # occupies.
                self._seen_filenames.discard(unique_md_name)
                await run_file_io(md_staging.unlink, True)
            else:
                unique_md_name, md_path = await self._link_commit_with_retry(uploads, Path(md_staged), unique_md_name)
                self._written_paths.append(md_path)
                md_virtual_path = uploads.upload_virtual_path(md_path.name)
                if self._sync_to_sandbox:
                    self._sync_targets.append((md_path, md_virtual_path))
                file_info["markdown_file"] = md_path.name
                file_info["markdown_path"] = str(self._uploads_dir / md_path.name)
                file_info["markdown_virtual_path"] = md_virtual_path
                file_info["markdown_artifact_url"] = uploads.upload_artifact_url(self._thread_id, md_path.name)
        return file_info

    async def finalize(self) -> None:
        """Make written files sandbox-readable, then sync to a non-mounted sandbox.

        Runs after every file of the session ingested successfully — exactly
        the ordinary endpoint's tail. Failures here propagate (the host files
        stay in place; the caller does not clean them up, matching ordinary
        upload behavior for sync-phase errors).
        """
        uploads = _uploads()
        # Uploaded files are created with 0o600 permissions (owner read/write
        # only). In Docker sandbox deployments the gateway writes as root but
        # the sandbox process runs as a non-root user (typically UID 1000).
        # Without group/other read bits the sandbox cannot access the files —
        # whether the uploads directory is bind-mounted into the container or
        # synced via sandbox.update_file. Always add group/other read bits so
        # every sandbox configuration can read the uploaded content.
        await run_file_io(uploads._make_uploaded_paths_sandbox_readable, self._written_paths)
        if self._sync_to_sandbox and self._sandbox is not None:
            for file_path, virtual_path in self._sync_targets:
                await run_file_io(uploads._sync_upload_to_sandbox, self._sandbox, file_path, virtual_path)

    async def cleanup_written(self) -> None:
        """Remove every file this session wrote (ordinary rejected-request cleanup)."""
        if not self._written_paths:
            return
        uploads = _uploads()
        await run_file_io(uploads._cleanup_uploaded_paths, self._written_paths)
        self._written_paths = []

    async def aclose(self) -> None:
        """Release the sandbox request lease (failures are logged, never raised)."""
        if self._sandbox_lease is None:
            return
        try:
            await self._sandbox_lease.release()
        except Exception:
            logger.warning(
                "Failed to release sandbox request lease after upload sync: %s",
                self._sandbox_lease.sandbox_id,
                exc_info=True,
            )