diff --git a/routes/personal_routes.py b/routes/personal_routes.py index a42615be7..53007e050 100644 --- a/routes/personal_routes.py +++ b/routes/personal_routes.py @@ -1,11 +1,13 @@ # routes/personal_routes.py """Routes for personal documents management.""" +import asyncio import os import logging import shutil import uuid from typing import Any, Dict, List, Tuple from fastapi import APIRouter, HTTPException, Query, Request, UploadFile, File, Depends +from fastapi.concurrency import run_in_threadpool from src.request_models import DirectoryRequest from core.constants import BASE_DIR, PERSONAL_DIR, PERSONAL_UPLOADS_DIR from src.rag_singleton import get_rag_manager @@ -18,7 +20,6 @@ UPLOADS_DIR = PERSONAL_UPLOADS_DIR logger = logging.getLogger(__name__) - def _personal_upload_dir_for_owner(owner: str | None, *, create: bool = True) -> str: """Return the per-owner upload directory used for direct RAG uploads.""" owner_segment = secure_filename((owner or "local").strip())[:80] or "local" @@ -141,6 +142,22 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): """ router = APIRouter(prefix="/api/personal") + # Serializes directory index jobs across requests. Indexing runs in the + # threadpool (#5558), so concurrent requests would otherwise run in parallel + # and race PersonalDocsManager's unsynchronized list mutations and file + # writes; before the threadpool move they serialized on the blocked event + # loop, so one-at-a-time is behavior parity. + # + # An asyncio.Lock acquired in the async handler BEFORE offloading: a waiting + # request parks on the event loop instead of pinning a threadpool worker (an + # earlier threading.Lock taken INSIDE the worker meant queued jobs held pool + # tokens while blocked, starving every other run_in_threadpool caller). + # add/remove/reload all take this lock, so their mutations never interleave. + # Per-router (not module-global) so each app binds it to its own event loop. + # Scope is the single process: multi-worker deployments would need a shared + # lock (out of scope for #5558). + _index_job_lock = asyncio.Lock() + def _rag(): """Get the current RAG manager, retrying init if needed.""" return get_rag_manager() @@ -172,8 +189,12 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): return {"files": files, "directories": directories} @router.post("/reload") - def api_personal_reload(owner: str = Depends(require_user), _admin: None = Depends(require_admin)): - personal_docs_manager.refresh_index() + async def api_personal_reload(owner: str = Depends(require_user), _admin: None = Depends(require_admin)): + # refresh_index() re-extracts text across every tracked directory — + # blocking work. Take the shared job lock (so it cannot race an add / + # remove) and run it off the event loop. + async with _index_job_lock: + await run_in_threadpool(personal_docs_manager.refresh_index) return {"ok": True, "count": len(personal_docs_manager.index)} @router.post("/add_directory") @@ -207,12 +228,26 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): # Use the RAGManager to index the directory rag = _rag() if rag: - result = rag.index_personal_documents(directory, owner=owner) - + def _index_directory(): + result = rag.index_personal_documents(directory, owner=owner) + if result["success"]: + # Also update the personal_docs_manager to track this + # directory. Kept inside the offloaded call: it triggers + # refresh_index(), which re-extracts text across tracked + # directories. + personal_docs_manager.add_directory(directory, index=False) + return result + + # Indexing walks, embeds, and stores the whole tree — minutes + # on a real directory. The handler is async, so calling it + # inline runs it on the event loop and every other request + # queues behind it until it finishes (#5558). Serialize on the + # async job lock BEFORE offloading so a queued request parks on + # the loop instead of pinning a threadpool worker. + async with _index_job_lock: + result = await run_in_threadpool(_index_directory) + if result["success"]: - # Also update the personal_docs_manager to track this directory - personal_docs_manager.add_directory(directory, index=False) - return { "success": True, "message": f"Successfully indexed {result['indexed_count']} chunks from {directory}", @@ -251,17 +286,25 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): logger.info(f"Removing directory from RAG: {directory}") - # Always remove from personal_docs_manager tracking - if hasattr(personal_docs_manager, 'remove_directory'): - personal_docs_manager.remove_directory(directory) - - # Remove from RAG vector store (best-effort) rag = _rag() - if rag: - try: - rag.remove_directory(directory) - except Exception as e: - logger.warning(f"RAG removal failed for directory {directory}: {e}") + + def _remove_directory(): + # Always remove from personal_docs_manager tracking. This + # mutates the same unsynchronized list/index an add job touches + # and re-extracts text (refresh_index), so it is blocking work. + if hasattr(personal_docs_manager, 'remove_directory'): + personal_docs_manager.remove_directory(directory) + # Remove from RAG vector store (best-effort). + if rag: + try: + rag.remove_directory(directory) + except Exception as e: + logger.warning(f"RAG removal failed for directory {directory}: {e}") + + # Same job lock as add/reload so remove cannot interleave with an + # in-flight add; offloaded off the event loop. + async with _index_job_lock: + await run_in_threadpool(_remove_directory) return { "success": True, diff --git a/tests/test_add_directory_event_loop.py b/tests/test_add_directory_event_loop.py new file mode 100644 index 000000000..7256d502a --- /dev/null +++ b/tests/test_add_directory_event_loop.py @@ -0,0 +1,253 @@ +"""Regression guard for #5558 — POST /api/personal/add_directory must not run +the indexing job on the event loop. + +The handler is ``async def`` but called ``rag.index_personal_documents`` +(os.walk + file reads + per-chunk embedding + Chroma inserts) inline, so +FastAPI ran the whole job on the event loop and every other request queued +behind it: indexing a real directory froze the UI and API for 25+ minutes. +``personal_docs_manager.add_directory`` sits in the same blocking section — it +triggers ``refresh_index()``, which re-extracts text across tracked dirs. + +These tests build the real router with fake managers and compare the thread +the indexing work runs on against the event loop's thread. +""" +import asyncio +import os +import threading + +os.environ.setdefault("DATABASE_URL", "sqlite:///:memory:") + +import httpx +from fastapi import FastAPI +from fastapi.testclient import TestClient + + +def _serialization_probe(): + """Shared counter proving two critical sections never overlap.""" + state = {"active": 0, "max_active": 0} + lock = threading.Lock() + + def enter(): + with lock: + state["active"] += 1 + state["max_active"] = max(state["max_active"], state["active"]) + + def leave(): + with lock: + state["active"] -= 1 + + return state, enter, leave + + +# Concurrency tests are `async def` (pyproject asyncio_mode="auto") and drive the +# ASGI app through httpx.ASGITransport + AsyncClient + asyncio.gather, NOT starlette +# TestClient + ThreadPoolExecutor: the job lock is an asyncio.Lock acquired in the +# async handler, and TestClient's portal-thread dispatch deadlocks against it (same +# reason test_notes_fail_closed_auth.py uses ASGITransport). asyncio.gather runs both +# requests on the test's own loop. +def _async_client(app): + return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://t") + +import routes.personal_routes as personal_routes +from core.middleware import require_admin +from src.auth_helpers import require_user + + +class _FakeRag: + def __init__(self, record): + self._record = record + + def index_personal_documents(self, directory, owner=None): + self._record["index_thread"] = threading.get_ident() + return {"success": True, "indexed_count": 3, "failed_count": 0} + + +class _FakeDocsManager: + def __init__(self, record): + self._record = record + self.index = [] + + def add_directory(self, directory, *, index=True, owner=None): + self._record["bookkeeping_thread"] = threading.get_ident() + self._record["bookkeeping_index_flag"] = index + + +def _build_app(tmp_path, monkeypatch, record): + monkeypatch.setattr(personal_routes, "PERSONAL_DIR", str(tmp_path)) + monkeypatch.setattr(personal_routes, "get_rag_manager", lambda: _FakeRag(record)) + + app = FastAPI() + app.include_router( + personal_routes.setup_personal_routes(_FakeDocsManager(record), None, True) + ) + app.dependency_overrides[require_user] = lambda: "tester" + app.dependency_overrides[require_admin] = lambda: None + + @app.get("/loop-thread") + async def loop_thread_probe(): + return {"thread": threading.get_ident()} + + return app + + +def test_indexing_runs_off_the_event_loop(tmp_path, monkeypatch): + record = {} + app = _build_app(tmp_path, monkeypatch, record) + target = tmp_path / "docs" + target.mkdir() + + # Context-manager client: one portal/event loop serves both requests, so + # the probe and the POST are guaranteed to see the same loop thread. + with TestClient(app) as client: + loop_thread = client.get("/loop-thread").json()["thread"] + resp = client.post( + "/api/personal/add_directory", json={"directory": str(target)} + ) + + assert resp.status_code == 200 + assert record["index_thread"] != loop_thread, ( + "index_personal_documents ran on the event loop thread — every other " + "request queues behind the indexing job (#5558)" + ) + assert record["bookkeeping_thread"] != loop_thread, ( + "personal_docs_manager.add_directory (refresh_index) ran on the event " + "loop thread" + ) + + +def test_response_and_bookkeeping_unchanged(tmp_path, monkeypatch): + record = {} + app = _build_app(tmp_path, monkeypatch, record) + target = tmp_path / "docs" + target.mkdir() + + client = TestClient(app) + resp = client.post("/api/personal/add_directory", json={"directory": str(target)}) + + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["indexed_count"] == 3 + assert body["failed_count"] == 0 + assert body["directory"] == os.path.realpath(str(target)) + assert record["bookkeeping_index_flag"] is False + + +async def test_concurrent_add_directory_requests_serialize_indexing(tmp_path, monkeypatch): + """Off-loop execution must not mean parallel index jobs: concurrent + requests would race PersonalDocsManager's unsynchronized list mutations + and file writes (save_directories/_save_excluded are plain open('w')).""" + import time + + state, enter, leave = _serialization_probe() + + def _slow_index(self, directory, owner=None): + enter(); time.sleep(0.2); leave() + return {"success": True, "indexed_count": 1, "failed_count": 0} + + monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index) + + record = {} + app = _build_app(tmp_path, monkeypatch, record) + for name in ("docs_a", "docs_b"): + (tmp_path / name).mkdir() + + async with _async_client(app) as ac: + results = await asyncio.gather( + ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_a")}), + ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_b")}), + ) + + assert all(r.status_code == 200 for r in results) + assert state["max_active"] == 1, ( + f"{state['max_active']} index jobs ran in parallel — concurrent " + "add_directory requests must serialize" + ) + + +def test_failed_indexing_still_returns_500(tmp_path, monkeypatch): + record = {} + app = _build_app(tmp_path, monkeypatch, record) + target = tmp_path / "docs" + target.mkdir() + + def _fail(directory, owner=None): + return {"success": False, "message": "boom"} + + monkeypatch.setattr(_FakeRag, "index_personal_documents", staticmethod(_fail)) + + client = TestClient(app) + resp = client.post("/api/personal/add_directory", json={"directory": str(target)}) + assert resp.status_code == 500 + assert "boom" in resp.json()["detail"] + + +async def test_add_and_remove_serialize(tmp_path, monkeypatch): + """#5634: remove must hold the SAME job lock as add. Otherwise a remove + running while an add job is in flight races PersonalDocsManager's + unsynchronized list/index mutations — the inconsistent state the PR's + 'add/remove are serialized' guarantee claims to prevent.""" + import time + + state, enter, leave = _serialization_probe() + + def _slow_index(self, directory, owner=None): + enter(); time.sleep(0.25); leave() + return {"success": True, "indexed_count": 1, "failed_count": 0} + + def _slow_remove(self, directory): + enter(); time.sleep(0.25); leave() + + monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index) + monkeypatch.setattr(_FakeDocsManager, "remove_directory", _slow_remove, raising=False) + + record = {} + app = _build_app(tmp_path, monkeypatch, record) + (tmp_path / "docs_a").mkdir() + (tmp_path / "docs_b").mkdir() + + async with _async_client(app) as ac: + results = await asyncio.gather( + ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_a")}), + ac.delete("/api/personal/remove_directory", params={"directory": str(tmp_path / "docs_b")}), + ) + + assert all(r.status_code == 200 for r in results) + assert state["max_active"] == 1, ( + f"{state['max_active']} add/remove critical sections overlapped — " + "remove must hold the same index job lock as add" + ) + + +async def test_reload_serializes_with_add(tmp_path, monkeypatch): + """#5634: POST /reload rebuilds the index via refresh_index(); it must hold + the same job lock so it cannot race an in-flight add job.""" + import time + + state, enter, leave = _serialization_probe() + + def _slow_index(self, directory, owner=None): + enter(); time.sleep(0.25); leave() + return {"success": True, "indexed_count": 1, "failed_count": 0} + + def _slow_refresh(self): + enter(); time.sleep(0.25); leave() + + monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index) + monkeypatch.setattr(_FakeDocsManager, "refresh_index", _slow_refresh, raising=False) + + record = {} + app = _build_app(tmp_path, monkeypatch, record) + (tmp_path / "docs_a").mkdir() + + async with _async_client(app) as ac: + results = await asyncio.gather( + ac.post("/api/personal/add_directory", json={"directory": str(tmp_path / "docs_a")}), + ac.post("/api/personal/reload"), + ) + + assert all(r.status_code == 200 for r in results) + assert state["max_active"] == 1, ( + f"{state['max_active']} add/reload critical sections overlapped — " + "reload must hold the same index job lock as add" + )