diff --git a/routes/personal_routes.py b/routes/personal_routes.py index 53007e050..c3201380a 100644 --- a/routes/personal_routes.py +++ b/routes/personal_routes.py @@ -328,10 +328,12 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): upload_dir = _personal_upload_dir_for_owner(user) - total_indexed = 0 total_failed = 0 - uploaded_files = [] + # Read the request bodies on the event loop — that part is genuine async + # I/O — then stage them so every blocking step happens in one offloaded + # critical section below. + staged: List[Tuple[str, str, str, bytes]] = [] for upload in files: try: file_path, stored_name, safe_name = _unique_personal_upload_path(upload_dir, upload.filename) @@ -340,46 +342,67 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): logger.warning(f"Rejected oversized personal upload: {upload.filename!r}") total_failed += 1 continue - with open(file_path, "wb") as f: - f.write(content_bytes) - - ext = os.path.splitext(safe_name)[1].lower() - if ext == ".pdf": - from src.personal_docs import extract_pdf_text - text = extract_pdf_text(file_path) - else: - text = content_bytes.decode("utf-8", errors="replace") - - if not text or not text.strip(): - total_failed += 1 - continue - - # Chunk and index - chunks = rag._split_into_chunks(text, chunk_size=500) - for i, chunk in enumerate(chunks): - metadata = { - "source": file_path, - "filename": safe_name, - "stored_filename": stored_name, - "directory": upload_dir, - "type": ext, - "chunk_id": i, - } - if user: - metadata["owner"] = user - if rag.add_document(chunk, metadata): - total_indexed += 1 - else: - total_failed += 1 - - uploaded_files.append(safe_name) + staged.append((file_path, stored_name, safe_name, content_bytes)) except Exception as e: - logger.error(f"Failed to upload/index {upload.filename}: {e}") + logger.error(f"Failed to read upload {upload.filename}: {e}") total_failed += 1 - # Track uploads directory - if uploaded_files and hasattr(personal_docs_manager, "add_directory"): - personal_docs_manager.add_directory(upload_dir, index=False) + def _index_uploads(): + indexed = 0 + failed = 0 + names = [] + for file_path, stored_name, safe_name, content_bytes in staged: + try: + with open(file_path, "wb") as f: + f.write(content_bytes) + + ext = os.path.splitext(safe_name)[1].lower() + if ext == ".pdf": + from src.personal_docs import extract_pdf_text + text = extract_pdf_text(file_path) + else: + text = content_bytes.decode("utf-8", errors="replace") + + if not text or not text.strip(): + failed += 1 + continue + + # Chunk and index + chunks = rag._split_into_chunks(text, chunk_size=500) + for i, chunk in enumerate(chunks): + metadata = { + "source": file_path, + "filename": safe_name, + "stored_filename": stored_name, + "directory": upload_dir, + "type": ext, + "chunk_id": i, + } + if user: + metadata["owner"] = user + if rag.add_document(chunk, metadata): + indexed += 1 + else: + failed += 1 + + names.append(safe_name) + except Exception as e: + logger.error(f"Failed to upload/index {safe_name}: {e}") + failed += 1 + + # Same transition, same lock: the tracking update must not land + # while another job is mid-write over the same state. + if names and hasattr(personal_docs_manager, "add_directory"): + personal_docs_manager.add_directory(upload_dir, index=False) + return indexed, failed, names + + # Chunking, embedding and the tracking update are blocking work over the + # same vector/tracking state add_directory mutates (#5634). Take the + # shared job lock BEFORE offloading so a queued request parks on the loop + # instead of pinning a threadpool worker, matching add_directory. + async with _index_job_lock: + total_indexed, indexed_failed, uploaded_files = await run_in_threadpool(_index_uploads) + total_failed += indexed_failed return { "success": True, @@ -392,38 +415,47 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available): async def delete_file_from_rag(filepath: str = Query(...), owner: str = Depends(require_user), _admin: None = Depends(require_admin)): """Delete a specific file from RAG index and optionally from disk.""" try: - # Remove chunks from RAG vector store (best-effort) - removed = 0 - rag = _rag() - if rag: - try: - removed = rag.delete_by_source(filepath) - except Exception as e: - logger.warning(f"RAG removal failed for {filepath}: {e}") + def _delete_file(): + # Remove chunks from RAG vector store (best-effort) + removed = 0 + rag = _rag() + if rag: + try: + removed = rag.delete_by_source(filepath) + except Exception as e: + logger.warning(f"RAG removal failed for {filepath}: {e}") - # Delete file from disk if it's in the caller's own uploads dir. - # Scope to the per-owner subdir, not the shared uploads root, so one - # admin can't delete another user's personal files by path. - deleted_from_disk = False - try: - abs_target = os.path.realpath(filepath) - base_abs = os.path.realpath(_personal_upload_dir_for_owner(owner, create=False)) - in_uploads = ( - abs_target == base_abs - or os.path.commonpath([abs_target, base_abs]) == base_abs - ) - except ValueError: - # commonpath raises on mixed drives / non-comparable paths - in_uploads = False - if in_uploads and abs_target != base_abs: + # Delete file from disk if it's in the caller's own uploads dir. + # Scope to the per-owner subdir, not the shared uploads root, so one + # admin can't delete another user's personal files by path. + deleted_from_disk = False try: - os.remove(abs_target) - deleted_from_disk = True - except FileNotFoundError: - pass # already gone — race with another request or cleanup + abs_target = os.path.realpath(filepath) + base_abs = os.path.realpath(_personal_upload_dir_for_owner(owner, create=False)) + in_uploads = ( + abs_target == base_abs + or os.path.commonpath([abs_target, base_abs]) == base_abs + ) + except ValueError: + # commonpath raises on mixed drives / non-comparable paths + in_uploads = False + if in_uploads and abs_target != base_abs: + try: + os.remove(abs_target) + deleted_from_disk = True + except FileNotFoundError: + pass # already gone — race with another request or cleanup - # Exclude the file from the listing (persists across restarts) - personal_docs_manager.exclude_file(filepath) + # Exclude the file from the listing (persists across restarts) + personal_docs_manager.exclude_file(filepath) + return removed, deleted_from_disk + + # Vector removal, the disk unlink and the exclusion write are one + # transition over the same state add_directory mutates (#5634), and + # all three block. Take the shared job lock BEFORE offloading, as + # add_directory does. + async with _index_job_lock: + removed, deleted_from_disk = await run_in_threadpool(_delete_file) return { "success": True, diff --git a/tests/test_add_directory_event_loop.py b/tests/test_add_directory_event_loop.py index 7256d502a..e677b3abb 100644 --- a/tests/test_add_directory_event_loop.py +++ b/tests/test_add_directory_event_loop.py @@ -61,6 +61,17 @@ class _FakeRag: self._record["index_thread"] = threading.get_ident() return {"success": True, "indexed_count": 3, "failed_count": 0} + def _split_into_chunks(self, text, chunk_size=500): + return [text] + + def add_document(self, chunk, metadata): + self._record["add_document_thread"] = threading.get_ident() + return True + + def delete_by_source(self, filepath): + self._record["delete_thread"] = threading.get_ident() + return 1 + class _FakeDocsManager: def __init__(self, record): @@ -71,6 +82,9 @@ class _FakeDocsManager: self._record["bookkeeping_thread"] = threading.get_ident() self._record["bookkeeping_index_flag"] = index + def exclude_file(self, filepath): + self._record["exclude_thread"] = threading.get_ident() + def _build_app(tmp_path, monkeypatch, record): monkeypatch.setattr(personal_routes, "PERSONAL_DIR", str(tmp_path)) @@ -219,6 +233,94 @@ async def test_add_and_remove_serialize(tmp_path, monkeypatch): ) +async def test_add_and_upload_serialize(tmp_path, monkeypatch): + """#5634 follow-up: POST /upload writes chunks into the vector store and then + calls personal_docs_manager.add_directory — the same vector/tracking state + add_directory mutates. It must hold the SAME job lock, or an upload landing + mid-add interleaves two writers over unsynchronized state.""" + 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_add_document(self, chunk, metadata): + self._record["add_document_thread"] = threading.get_ident() + enter(); time.sleep(0.25); leave() + return True + + monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index) + monkeypatch.setattr(_FakeRag, "add_document", _slow_add_document) + + record = {} + app = _build_app(tmp_path, monkeypatch, record) + monkeypatch.setattr(personal_routes, "UPLOADS_DIR", str(tmp_path / "uploads")) + monkeypatch.setattr(personal_routes, "require_privilege", lambda request, key: "tester") + (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/upload", files={"files": ("a.txt", b"hello world", "text/plain")}), + ) + + assert all(r.status_code == 200 for r in results) + # The test coroutine runs on the event loop, so this IS the loop thread. + assert record["add_document_thread"] != threading.get_ident(), ( + "rag.add_document ran on the event loop thread — chunk writes block " + "every other request for the duration of the upload" + ) + assert state["max_active"] == 1, ( + f"{state['max_active']} add/upload critical sections overlapped — " + "upload must hold the same index job lock as add" + ) + + +async def test_add_and_delete_file_serialize(tmp_path, monkeypatch): + """#5634 follow-up: DELETE /file removes chunks from the vector store and + calls personal_docs_manager.exclude_file. Both mutate state add_directory + also touches, so the delete must hold the SAME job lock as add.""" + 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_delete(self, filepath): + self._record["delete_thread"] = threading.get_ident() + enter(); time.sleep(0.25); leave() + return 1 + + monkeypatch.setattr(_FakeRag, "index_personal_documents", _slow_index) + monkeypatch.setattr(_FakeRag, "delete_by_source", _slow_delete) + + record = {} + app = _build_app(tmp_path, monkeypatch, record) + monkeypatch.setattr(personal_routes, "UPLOADS_DIR", str(tmp_path / "uploads")) + (tmp_path / "docs_a").mkdir() + doomed = tmp_path / "doomed.txt" + doomed.write_text("bye") + + 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/file", params={"filepath": str(doomed)}), + ) + + assert all(r.status_code == 200 for r in results) + assert record["delete_thread"] != threading.get_ident(), ( + "rag.delete_by_source ran on the event loop thread" + ) + assert state["max_active"] == 1, ( + f"{state['max_active']} add/delete critical sections overlapped — " + "delete 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."""