mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-08-24 01:44:18 +02:00
Compare commits
5 commits
dev
...
fix/add-di
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
96c88c27c8 | ||
|
|
6edd771cc9 |
||
|
|
938251000b | ||
|
|
e222e92153 | ||
|
|
b91f48f50a |
2 changed files with 560 additions and 95 deletions
|
|
@ -1,11 +1,13 @@
|
||||||
# routes/personal_routes.py
|
# routes/personal_routes.py
|
||||||
"""Routes for personal documents management."""
|
"""Routes for personal documents management."""
|
||||||
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
import uuid
|
import uuid
|
||||||
from typing import Any, Dict, List, Tuple
|
from typing import Any, Dict, List, Tuple
|
||||||
from fastapi import APIRouter, HTTPException, Query, Request, UploadFile, File, Depends
|
from fastapi import APIRouter, HTTPException, Query, Request, UploadFile, File, Depends
|
||||||
|
from fastapi.concurrency import run_in_threadpool
|
||||||
from src.request_models import DirectoryRequest
|
from src.request_models import DirectoryRequest
|
||||||
from core.constants import BASE_DIR, PERSONAL_DIR, PERSONAL_UPLOADS_DIR
|
from core.constants import BASE_DIR, PERSONAL_DIR, PERSONAL_UPLOADS_DIR
|
||||||
from src.rag_singleton import get_rag_manager
|
from src.rag_singleton import get_rag_manager
|
||||||
|
|
@ -18,7 +20,6 @@ UPLOADS_DIR = PERSONAL_UPLOADS_DIR
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
def _personal_upload_dir_for_owner(owner: str | None, *, create: bool = True) -> str:
|
def _personal_upload_dir_for_owner(owner: str | None, *, create: bool = True) -> str:
|
||||||
"""Return the per-owner upload directory used for direct RAG uploads."""
|
"""Return the per-owner upload directory used for direct RAG uploads."""
|
||||||
owner_segment = secure_filename((owner or "local").strip())[:80] or "local"
|
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")
|
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():
|
def _rag():
|
||||||
"""Get the current RAG manager, retrying init if needed."""
|
"""Get the current RAG manager, retrying init if needed."""
|
||||||
return get_rag_manager()
|
return get_rag_manager()
|
||||||
|
|
@ -172,8 +189,12 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
|
||||||
return {"files": files, "directories": directories}
|
return {"files": files, "directories": directories}
|
||||||
|
|
||||||
@router.post("/reload")
|
@router.post("/reload")
|
||||||
def api_personal_reload(owner: str = Depends(require_user), _admin: None = Depends(require_admin)):
|
async def api_personal_reload(owner: str = Depends(require_user), _admin: None = Depends(require_admin)):
|
||||||
personal_docs_manager.refresh_index()
|
# 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)}
|
return {"ok": True, "count": len(personal_docs_manager.index)}
|
||||||
|
|
||||||
@router.post("/add_directory")
|
@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
|
# Use the RAGManager to index the directory
|
||||||
rag = _rag()
|
rag = _rag()
|
||||||
if 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"]:
|
if result["success"]:
|
||||||
# Also update the personal_docs_manager to track this directory
|
|
||||||
personal_docs_manager.add_directory(directory, index=False)
|
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"success": True,
|
"success": True,
|
||||||
"message": f"Successfully indexed {result['indexed_count']} chunks from {directory}",
|
"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}")
|
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()
|
rag = _rag()
|
||||||
if rag:
|
|
||||||
try:
|
def _remove_directory():
|
||||||
rag.remove_directory(directory)
|
# Always remove from personal_docs_manager tracking. This
|
||||||
except Exception as e:
|
# mutates the same unsynchronized list/index an add job touches
|
||||||
logger.warning(f"RAG removal failed for directory {directory}: {e}")
|
# 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 {
|
return {
|
||||||
"success": True,
|
"success": True,
|
||||||
|
|
@ -289,54 +332,73 @@ def setup_personal_routes(personal_docs_manager, rag_manager, rag_available):
|
||||||
total_failed = 0
|
total_failed = 0
|
||||||
uploaded_files = []
|
uploaded_files = []
|
||||||
|
|
||||||
for upload in files:
|
# Chunking, embedding and the tracking update are blocking work over the
|
||||||
try:
|
# same vector/tracking state add_directory mutates (#5634). Take the
|
||||||
file_path, stored_name, safe_name = _unique_personal_upload_path(upload_dir, upload.filename)
|
# shared job lock BEFORE offloading so a queued request parks on the loop
|
||||||
content_bytes = await upload.read(PERSONAL_UPLOAD_MAX_BYTES + 1)
|
# instead of pinning a threadpool worker, matching add_directory.
|
||||||
if len(content_bytes) > PERSONAL_UPLOAD_MAX_BYTES:
|
# Read and process one capped payload at a time so a multi-file request
|
||||||
logger.warning(f"Rejected oversized personal upload: {upload.filename!r}")
|
# cannot retain len(files) * PERSONAL_UPLOAD_MAX_BYTES in memory.
|
||||||
total_failed += 1
|
async with _index_job_lock:
|
||||||
continue
|
for upload in files:
|
||||||
with open(file_path, "wb") as f:
|
try:
|
||||||
f.write(content_bytes)
|
file_path, stored_name, safe_name = _unique_personal_upload_path(
|
||||||
|
upload_dir, upload.filename
|
||||||
ext = os.path.splitext(safe_name)[1].lower()
|
)
|
||||||
if ext == ".pdf":
|
content_bytes = await upload.read(PERSONAL_UPLOAD_MAX_BYTES + 1)
|
||||||
from src.personal_docs import extract_pdf_text
|
if len(content_bytes) > PERSONAL_UPLOAD_MAX_BYTES:
|
||||||
text = extract_pdf_text(file_path)
|
logger.warning(f"Rejected oversized personal upload: {upload.filename!r}")
|
||||||
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
|
total_failed += 1
|
||||||
|
continue
|
||||||
|
|
||||||
uploaded_files.append(safe_name)
|
def _index_upload():
|
||||||
except Exception as e:
|
with open(file_path, "wb") as f:
|
||||||
logger.error(f"Failed to upload/index {upload.filename}: {e}")
|
f.write(content_bytes)
|
||||||
total_failed += 1
|
|
||||||
|
|
||||||
# Track uploads directory
|
ext = os.path.splitext(safe_name)[1].lower()
|
||||||
if uploaded_files and hasattr(personal_docs_manager, "add_directory"):
|
if ext == ".pdf":
|
||||||
personal_docs_manager.add_directory(upload_dir, index=False)
|
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():
|
||||||
|
return 0, 1, None
|
||||||
|
|
||||||
|
indexed = 0
|
||||||
|
failed = 0
|
||||||
|
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
|
||||||
|
return indexed, failed, safe_name
|
||||||
|
|
||||||
|
indexed, failed, uploaded_name = await run_in_threadpool(_index_upload)
|
||||||
|
total_indexed += indexed
|
||||||
|
total_failed += failed
|
||||||
|
if uploaded_name:
|
||||||
|
uploaded_files.append(uploaded_name)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to upload/index {upload.filename}: {e}")
|
||||||
|
total_failed += 1
|
||||||
|
|
||||||
|
# Same transition, same lock: the tracking update must not land
|
||||||
|
# while another job is mid-write over the same state.
|
||||||
|
if uploaded_files and hasattr(personal_docs_manager, "add_directory"):
|
||||||
|
await run_in_threadpool(
|
||||||
|
personal_docs_manager.add_directory, upload_dir, index=False
|
||||||
|
)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"success": True,
|
"success": True,
|
||||||
|
|
@ -349,38 +411,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)):
|
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."""
|
"""Delete a specific file from RAG index and optionally from disk."""
|
||||||
try:
|
try:
|
||||||
# Remove chunks from RAG vector store (best-effort)
|
def _delete_file():
|
||||||
removed = 0
|
# Remove chunks from RAG vector store (best-effort)
|
||||||
rag = _rag()
|
removed = 0
|
||||||
if rag:
|
rag = _rag()
|
||||||
try:
|
if rag:
|
||||||
removed = rag.delete_by_source(filepath)
|
try:
|
||||||
except Exception as e:
|
removed = rag.delete_by_source(filepath)
|
||||||
logger.warning(f"RAG removal failed for {filepath}: {e}")
|
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.
|
# 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
|
# Scope to the per-owner subdir, not the shared uploads root, so one
|
||||||
# admin can't delete another user's personal files by path.
|
# admin can't delete another user's personal files by path.
|
||||||
deleted_from_disk = False
|
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:
|
|
||||||
try:
|
try:
|
||||||
os.remove(abs_target)
|
abs_target = os.path.realpath(filepath)
|
||||||
deleted_from_disk = True
|
base_abs = os.path.realpath(_personal_upload_dir_for_owner(owner, create=False))
|
||||||
except FileNotFoundError:
|
in_uploads = (
|
||||||
pass # already gone — race with another request or cleanup
|
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)
|
# Exclude the file from the listing (persists across restarts)
|
||||||
personal_docs_manager.exclude_file(filepath)
|
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 {
|
return {
|
||||||
"success": True,
|
"success": True,
|
||||||
|
|
|
||||||
394
tests/test_add_directory_event_loop.py
Normal file
394
tests/test_add_directory_event_loop.py
Normal file
|
|
@ -0,0 +1,394 @@
|
||||||
|
"""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}
|
||||||
|
|
||||||
|
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):
|
||||||
|
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 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))
|
||||||
|
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_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_upload_processes_each_payload_before_reading_the_next(tmp_path, monkeypatch):
|
||||||
|
"""A multi-file upload must retain at most one capped payload at a time."""
|
||||||
|
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||||
|
|
||||||
|
reads = []
|
||||||
|
original_read = StarletteUploadFile.read
|
||||||
|
|
||||||
|
async def _recording_read(upload, size=-1):
|
||||||
|
reads.append(upload.filename)
|
||||||
|
return await original_read(upload, size)
|
||||||
|
|
||||||
|
def _record_first_index(self, chunk, metadata):
|
||||||
|
self._record.setdefault("reads_at_first_index", len(reads))
|
||||||
|
return True
|
||||||
|
|
||||||
|
monkeypatch.setattr(StarletteUploadFile, "read", _recording_read)
|
||||||
|
monkeypatch.setattr(_FakeRag, "add_document", _record_first_index)
|
||||||
|
|
||||||
|
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")
|
||||||
|
|
||||||
|
files = [
|
||||||
|
("files", ("a.txt", b"alpha", "text/plain")),
|
||||||
|
("files", ("b.txt", b"bravo", "text/plain")),
|
||||||
|
("files", ("c.txt", b"charlie", "text/plain")),
|
||||||
|
]
|
||||||
|
async with _async_client(app) as ac:
|
||||||
|
response = await ac.post("/api/personal/upload", files=files)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert response.json()["uploaded"] == ["a.txt", "b.txt", "c.txt"]
|
||||||
|
assert reads == ["a.txt", "b.txt", "c.txt"]
|
||||||
|
assert record["reads_at_first_index"] == 1, (
|
||||||
|
"all upload bodies were retained before worker processing began"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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."""
|
||||||
|
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"
|
||||||
|
)
|
||||||
Loading…
Add table
Add a link
Reference in a new issue