diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e6f4bfbf..cec318a0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,6 +31,7 @@ jobs: python -m pip install --upgrade pip pip install -r requirements.txt pip install -r test_requirements.txt + pip install -r requirements.extraction.txt - name: Run unit tests env: diff --git a/EXTRACTION.md b/EXTRACTION.md new file mode 100644 index 00000000..5678aa37 --- /dev/null +++ b/EXTRACTION.md @@ -0,0 +1,56 @@ +# Opt-in document extraction contract + +`POST /v1/extract` implements a **document-v1** profile for DOCX only. It does not +replace `/text`, alter ingestion/chunking, or invoke embeddings, OCR or the vector +store. The route is disabled by default and requires a verified `JWT_SECRET` +token with an `id` claim when enabled. The main RAG application **still initializes +its vector store and embeddings at startup**. A parsing-only deployment remains a +separate migration. + +Install the pinned optional engine in a custom image or Python environment, +then opt in explicitly before starting (or restarting) the API. With the flag +off, the route is not registered and does not parse multipart bodies: + +```sh +pip install -r requirements.extraction.txt +export RAG_EXTRACTION_API_ENABLED=true +``` + +Send multipart `file` and `profile=document-v1`. Use DOCX MIME or a `.docx` +filename with a generic MIME. Success includes `text` (Markdown), `format`, +`profile`, `completeness` (`complete`/`partial`), `may_omit_content`, +`pages_needing_ocr` (empty until PDF support), `truncated` (always false), and +`parser: {name, version}`. An embedded image marks the DOCX as *partial*. Never +use partial text as proof of full content inspection. `complete` means the +supported conversion finished without *known* omitted image entries, not that +all information in the source is provably inspectable. No hosted OCR or +second-parser fallback runs inside this endpoint. + +Failures use `detail.code`, not the native exception message: + +| Status | Codes | Action | +|---|---|---| +| 400/415 | `UNSUPPORTED_PROFILE`, `UNSUPPORTED_DOCUMENT_TYPE` | Select a supported profile/type | +| 401/404 | `EXTRACTION_AUTH_REQUIRED`, `EXTRACTION_DISABLED` | Authenticate/opt in | +| 413 | `PARSER_INPUT_LIMIT`, `PARSER_OUTPUT_LIMIT`, `ZIP_BOMB` | Hard refusal; never send the same bytes to another parser | +| 422 | `ARCHIVE_INVALID`, `NO_DOCUMENT_TEXT`, `PARSE_FAILED` | Unusable archive or empty/unconvertible document | +| 429 | `CONCURRENCY_LIMIT` | Retry later; not a reason to invoke paid OCR | +| 503/504 | `PARSER_UNAVAILABLE`, `PARSER_CRASH`, `PARSER_TIMEOUT` | Retry or fix the service | + +The route checks the input limit of 15 MiB while staging the upload; +serialized output is capped at 15 MiB before IPC. Starlette may have already +spooled a multipart upload before the route runs: configure an upstream HTTP +body-size limit as well for internet-facing deployments. The child checks +*actual decompressed* ZIP entry bytes: at most +25 MiB per entry, 100 MiB in total and 4,096 entries. Defaults are two active +parses and six queued per API process. Set `RAG_EXTRACTION_CONCURRENT`, +`RAG_EXTRACTION_QUEUED` and `RAG_EXTRACTION_TIMEOUT_SECONDS` to tune admission +and the overall 30-second default deadline (queue wait, upload staging, parse). +On cancellation/timeout the child is killed and reaped before its temp file is +removed and its slot is reused. + +This is the first **service-side** slice. Existing LibreChat local parsing and +RAG `/text` behavior remain in place until cross-service tests establish policy, +authorization, preview, failure and compatibility behavior for each consumer. +The real DOCX test fixture is copied from Marco's LibreChat AnyDoc PR #14701 at +`fb7bbcd9cf75f4f78ecbd5a8780685c481600be2`. diff --git a/app/models.py b/app/models.py index 57c01130..684b859c 100644 --- a/app/models.py +++ b/app/models.py @@ -2,7 +2,23 @@ import hashlib from enum import Enum from pydantic import BaseModel -from typing import Optional, List +from typing import Optional, List, Literal + + +class ParserProvenance(BaseModel): + name: Literal["anydoc"] + version: str + + +class ExtractionResult(BaseModel): + profile: Literal["document-v1"] + text: str + format: Literal["markdown"] + completeness: Literal["complete", "partial"] + may_omit_content: bool + pages_needing_ocr: List[int] + truncated: bool + parser: ParserProvenance class DocumentResponse(BaseModel): diff --git a/app/routes/extraction_routes.py b/app/routes/extraction_routes.py new file mode 100644 index 00000000..1082a6ad --- /dev/null +++ b/app/routes/extraction_routes.py @@ -0,0 +1,124 @@ +"""Versioned document extraction. No embeddings, vector writes, or OCR calls.""" + +import asyncio +import math +import os +import tempfile +from pathlib import Path + +import aiofiles +from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile + +from app.config import RAG_UPLOAD_DIR, logger +from app.models import ExtractionResult +from app.services.extraction import ( + ExtractionAdmission, + ExtractionBusy, + ExtractionFailure, + run_worker, +) + +router = APIRouter(prefix="/v1") +DOCX_TYPE = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" +MAX_INPUT_BYTES = 15 * 1024 * 1024 +_DEFAULT_TIMEOUT = 30.0 +_admission: ExtractionAdmission | None = None + + +def _error(code: str, status_code: int) -> HTTPException: + return HTTPException(status_code=status_code, detail={"code": code}) + + +def _get_admission() -> ExtractionAdmission: + global _admission + if _admission is None: + _admission = ExtractionAdmission( + concurrent=int(os.getenv("RAG_EXTRACTION_CONCURRENT", "2")), + queued=int(os.getenv("RAG_EXTRACTION_QUEUED", "6")), + ) + return _admission + + +async def _save_bounded(file: UploadFile, path: Path) -> None: + size = 0 + async with aiofiles.open(path, "wb") as output: + while chunk := await file.read(64 * 1024): + size += len(chunk) + if size > MAX_INPUT_BYTES: + raise _error("PARSER_INPUT_LIMIT", 413) + await output.write(chunk) + + +@router.post("/extract", response_model=ExtractionResult) +async def extract_document( + request: Request, + file: UploadFile = File(...), + profile: str = Form(...), +) -> ExtractionResult: + # Existing /text remains unchanged. An operator must explicitly enable + # and install this separate profile before moving any LibreChat caller. + if os.getenv("RAG_EXTRACTION_API_ENABLED", "false").lower() not in { + "1", + "true", + "yes", + "on", + }: + raise _error("EXTRACTION_DISABLED", 404) + # Legacy RAG deployments may run without auth; expensive extraction is + # never allowed anonymously, even when those older routes are public. + if not os.getenv("JWT_SECRET") or not getattr(request.state, "user", {}).get("id"): + raise _error("EXTRACTION_AUTH_REQUIRED", 401) + if profile != "document-v1": + raise _error("UNSUPPORTED_PROFILE", 400) + content_type = (file.content_type or "").split(";")[0].strip().lower() + extension = Path(file.filename or "").suffix.lower() + if content_type == "application/pdf" or not ( + content_type == DOCX_TYPE + or ( + content_type in {"application/octet-stream", "binary/octet-stream", ""} + and extension == ".docx" + ) + ): + raise _error("UNSUPPORTED_DOCUMENT_TYPE", 415) + if file.size is not None and file.size > MAX_INPUT_BYTES: + raise _error("PARSER_INPUT_LIMIT", 413) + + try: + admission = _get_admission() + timeout = float( + os.getenv("RAG_EXTRACTION_TIMEOUT_SECONDS", str(_DEFAULT_TIMEOUT)) + ) + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("Invalid extraction timeout") + async with asyncio.timeout(timeout): + async with admission.slot(): + fd, filename = tempfile.mkstemp( + prefix="rag-extract-", suffix=".docx", dir=RAG_UPLOAD_DIR + ) + os.close(fd) + path = Path(filename) + try: + await _save_bounded(file, path) + return await run_worker(path) + finally: + path.unlink(missing_ok=True) + except ExtractionBusy: + raise _error("CONCURRENCY_LIMIT", 429) + except ExtractionFailure as exc: + status_code = { + "ZIP_BOMB": 413, + "ARCHIVE_INVALID": 422, + "PARSER_OUTPUT_LIMIT": 413, + "NO_DOCUMENT_TEXT": 422, + "PARSE_FAILED": 422, + "PARSER_UNAVAILABLE": 503, + "PARSER_CRASH": 503, + }[exc.code] + raise _error(exc.code, status_code) + except TimeoutError: + raise _error("PARSER_TIMEOUT", 504) + except (OSError, ValueError) as exc: + logger.error( + "Extraction infrastructure unavailable | error=%s", type(exc).__name__ + ) + raise _error("PARSER_UNAVAILABLE", 503) diff --git a/app/services/extraction.py b/app/services/extraction.py new file mode 100644 index 00000000..06148362 --- /dev/null +++ b/app/services/extraction.py @@ -0,0 +1,100 @@ +"""Bounded, cancellable process boundary for opt-in document extraction.""" + +import asyncio +import json +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncIterator + +from app.models import ExtractionResult + +MAX_IPC_BYTES = 15 * 1024 * 1024 + + +class ExtractionBusy(Exception): + pass + + +class ExtractionFailure(Exception): + def __init__(self, code: str): + self.code = code + + +class ExtractionAdmission: + """One per serving process: bound uploads waiting and native children running.""" + + def __init__(self, concurrent: int = 2, queued: int = 6): + if concurrent < 1 or queued < 0: + raise ValueError("Invalid extraction admission limits") + self._capacity = concurrent + queued + self._pending = 0 + self._slots = asyncio.Semaphore(concurrent) + + @asynccontextmanager + async def slot(self) -> AsyncIterator[None]: + # The serving process has one event loop. No await separates the check + # and increment, so concurrent requests cannot exceed the queue limit. + if self._pending >= self._capacity: + raise ExtractionBusy() + self._pending += 1 + acquired = False + try: + await self._slots.acquire() + acquired = True + yield + finally: + self._pending -= 1 + if acquired: + self._slots.release() + + +def _command(path: Path) -> tuple[str, ...]: + return (sys.executable, "-m", "app.services.extraction_worker", str(path)) + + +async def run_worker(path: Path) -> ExtractionResult: + """Read a bounded child response; always reap a child before releasing its slot.""" + process = await asyncio.create_subprocess_exec( + *_command(path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + cwd=str(Path(__file__).resolve().parents[2]), + ) + try: + output = bytearray() + while chunk := await process.stdout.read(64 * 1024): + output.extend(chunk) + if len(output) > MAX_IPC_BYTES: + raise ExtractionFailure("PARSER_OUTPUT_LIMIT") + await process.wait() + except BaseException: + if process.returncode is None: + try: + process.kill() + except ProcessLookupError: + pass # The child exited while cancellation was being delivered. + # Draining and reaping are necessary before the temporary file can be + # removed and admission can be granted to the next upload. + await process.communicate() + raise + if process.returncode != 0: + raise ExtractionFailure("PARSER_CRASH") + try: + message = json.loads(output) + if message.get("ok") is False: + code = message["code"] + if code in { + "ZIP_BOMB", + "ARCHIVE_INVALID", + "PARSER_UNAVAILABLE", + "PARSER_OUTPUT_LIMIT", + "NO_DOCUMENT_TEXT", + "PARSE_FAILED", + }: + raise ExtractionFailure(code) + return ExtractionResult.model_validate(message["result"]) + except ExtractionFailure: + raise + except (AttributeError, KeyError, TypeError, ValueError) as exc: + raise ExtractionFailure("PARSER_CRASH") from exc diff --git a/app/services/extraction_worker.py b/app/services/extraction_worker.py new file mode 100644 index 00000000..8bd827be --- /dev/null +++ b/app/services/extraction_worker.py @@ -0,0 +1,136 @@ +"""Isolated document parsing. Run only as ``python -m app.services.extraction_worker``. + +The web process does not import native parsing bindings. Child exit, crash, and +SIGKILL cannot terminate the API process or leave its event loop blocked. +""" + +import json +import sys +import zipfile +from importlib.metadata import version +from pathlib import Path + +MAX_ARCHIVE_ENTRIES = 4096 +MAX_ENTRY_BYTES = 25 * 1024 * 1024 +MAX_TOTAL_BYTES = 100 * 1024 * 1024 +MAX_OUTPUT_BYTES = 15 * 1024 * 1024 +IMAGE_EXTENSIONS = ( + ".jpg", + ".jpeg", + ".png", + ".gif", + ".tif", + ".tiff", + ".bmp", + ".webp", + ".jp2", + ".jpx", + ".avif", + ".heic", + ".heif", + ".emf", + ".wmf", + ".svg", +) + + +class ExtractionRefusal(Exception): + def __init__(self, code: str): + self.code = code + + +def inspect_docx(path: Path) -> bool: + """Validate *actually inflated* archive bytes before handing any to AnyDoc. + + Reading in 64-KiB chunks catches false central-directory sizes without + keeping decompressed entries in memory. This is deliberately separate from + the text-output limit, since even an empty parse can inflate a zip bomb. + """ + try: + with zipfile.ZipFile(path) as archive: + entries = archive.infolist() + if len(entries) > MAX_ARCHIVE_ENTRIES: + raise ExtractionRefusal("ZIP_BOMB") + names = {entry.filename for entry in entries} + if not {"[Content_Types].xml", "word/document.xml"} <= names: + raise ExtractionRefusal("ARCHIVE_INVALID") + total = 0 + may_omit_content = False + for entry in entries: + if entry.is_dir(): + continue + if ( + entry.file_size > MAX_ENTRY_BYTES + or total + entry.file_size > MAX_TOTAL_BYTES + ): + raise ExtractionRefusal("ZIP_BOMB") + name = entry.filename.lower() + if not name.startswith(("docprops/", "thumbnails/")) and name.endswith( + IMAGE_EXTENSIONS + ): + may_omit_content = True + entry_bytes = 0 + with archive.open(entry) as stream: + while chunk := stream.read(64 * 1024): + entry_bytes += len(chunk) + total += len(chunk) + if entry_bytes > MAX_ENTRY_BYTES or total > MAX_TOTAL_BYTES: + raise ExtractionRefusal("ZIP_BOMB") + return may_omit_content + except ( + zipfile.BadZipFile, + zipfile.LargeZipFile, + RuntimeError, + EOFError, + NotImplementedError, + OSError, + ) as exc: + raise ExtractionRefusal("ARCHIVE_INVALID") from exc + + +def extract(path: Path) -> dict: + may_omit_content = inspect_docx(path) + try: + import anydoc + + # The MIME/extension are never passed to the binding. DOCX identity was + # checked in the archive above; explicit format avoids filename fallback. + text = anydoc.to_markdown_bytes(path.read_bytes(), "docx") + except ImportError as exc: + raise ExtractionRefusal("PARSER_UNAVAILABLE") from exc + except Exception as exc: + # Native error messages can echo source content. Never send them to clients. + raise ExtractionRefusal("PARSE_FAILED") from exc + if not isinstance(text, str) or not text.strip(): + raise ExtractionRefusal("NO_DOCUMENT_TEXT") + if len(text.encode("utf-8")) > MAX_OUTPUT_BYTES: + raise ExtractionRefusal("PARSER_OUTPUT_LIMIT") + return { + "profile": "document-v1", + "text": text, + "format": "markdown", + "completeness": "partial" if may_omit_content else "complete", + "may_omit_content": may_omit_content, + "pages_needing_ocr": [], + "truncated": False, + "parser": {"name": "anydoc", "version": version("firecrawl-anydoc")}, + } + + +def main() -> None: + if len(sys.argv) != 2: + raise SystemExit(2) + try: + payload = {"ok": True, "result": extract(Path(sys.argv[1]))} + except ExtractionRefusal as exc: + payload = {"ok": False, "code": exc.code} + except Exception: + payload = {"ok": False, "code": "PARSE_FAILED"} + serialized = json.dumps(payload, ensure_ascii=False) + if len(serialized.encode("utf-8")) > MAX_OUTPUT_BYTES: + serialized = json.dumps({"ok": False, "code": "PARSER_OUTPUT_LIMIT"}) + sys.stdout.write(serialized) + + +if __name__ == "__main__": + main() diff --git a/main.py b/main.py index e300951a..714f5752 100644 --- a/main.py +++ b/main.py @@ -23,7 +23,7 @@ vector_store, ) from app.middleware import security_middleware -from app.routes import document_routes, pgvector_routes +from app.routes import document_routes, extraction_routes, pgvector_routes from app.services.database import PSQLDatabase, ensure_vector_indexes from app.services.vector_store.factory import close_vector_store_connections @@ -90,6 +90,13 @@ async def lifespan(app: FastAPI): # Include routers app.include_router(document_routes.router) +if os.getenv("RAG_EXTRACTION_API_ENABLED", "false").lower() in { + "1", + "true", + "yes", + "on", +}: + app.include_router(extraction_routes.router) if debug_mode: app.include_router(router=pgvector_routes.router) diff --git a/requirements.extraction.txt b/requirements.extraction.txt new file mode 100644 index 00000000..e8276a3a --- /dev/null +++ b/requirements.extraction.txt @@ -0,0 +1,3 @@ +# Optional document-v1 extraction engine; install only when enabling /v1/extract. +# Match LibreChat AnyDoc PR #14701 (Node @firecrawl/anydoc 0.1.3) for corpus parity. +firecrawl-anydoc==0.1.3 diff --git a/tests/fixtures/structured.docx b/tests/fixtures/structured.docx new file mode 100644 index 00000000..eff663f8 Binary files /dev/null and b/tests/fixtures/structured.docx differ diff --git a/tests/test_extraction_api.py b/tests/test_extraction_api.py new file mode 100644 index 00000000..59dc1592 --- /dev/null +++ b/tests/test_extraction_api.py @@ -0,0 +1,393 @@ +"""Contract tests for the opt-in extraction slice; uses the pinned native wheel.""" + +import asyncio +import io +import os +import subprocess +import sys +import zipfile +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import httpx +import jwt +import pytest +from fastapi import FastAPI, HTTPException, UploadFile + +from app.middleware import security_middleware +from app.routes import extraction_routes +from app.services import extraction, extraction_worker +from main import app as main_app + +app = FastAPI() +app.middleware("http")(security_middleware) +app.include_router(extraction_routes.router) + +FIXTURE = Path(__file__).parent / "fixtures" / "structured.docx" +DOCX_TYPE = extraction_routes.DOCX_TYPE + + +@pytest.fixture +def configured(monkeypatch, tmp_path): + monkeypatch.setenv("RAG_EXTRACTION_API_ENABLED", "true") + monkeypatch.setenv("JWT_SECRET", "a-test-key-that-is-at-least-32-bytes-long") + monkeypatch.setattr(extraction_routes, "RAG_UPLOAD_DIR", str(tmp_path)) + admission = extraction.ExtractionAdmission() + monkeypatch.setattr(extraction_routes, "_admission", admission) + with ThreadPoolExecutor(max_workers=2) as pool: + monkeypatch.setattr(main_app.state, "thread_pool", pool, raising=False) + yield tmp_path, admission + + +@pytest.fixture +def headers(configured): + token = jwt.encode({"id": "owner"}, os.environ["JWT_SECRET"], algorithm="HS256") + return {"Authorization": f"Bearer {token}"} + + +async def post( + file_bytes, headers, name="report.docx", mime=DOCX_TYPE, profile="document-v1" +): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as client: + return await client.post( + "/v1/extract", + headers=headers, + data={"profile": profile}, + files={"file": (name, file_bytes, mime)}, + ) + + +def with_entry(original: bytes, name: str, value: bytes) -> bytes: + result = io.BytesIO() + with zipfile.ZipFile(io.BytesIO(original)) as source, zipfile.ZipFile( + result, "w", zipfile.ZIP_DEFLATED + ) as target: + for entry in source.infolist(): + if entry.filename != name: + target.writestr(entry, source.read(entry)) + target.writestr(name, value) + return result.getvalue() + + +async def test_real_docx_from_marcos_pr_returns_markdown_and_provenance( + headers, configured +): + response = await post( + FIXTURE.read_bytes(), headers, name="renamed.csv", mime=DOCX_TYPE + ) + assert response.status_code == 200, response.text + payload = response.json() + assert payload == { + "profile": "document-v1", + "format": "markdown", + "text": ( + "# Quarterly Report\n\nThis document summarizes the results for the period.\n\n" + "## Regional Totals\n\n| | | |\n| --- | --- | --- |\n" + "| Region | Units | Revenue |\n| North | 1200 | 48000 |\n" + "| South | 950 | 38000 |\n| East | 1430 | 57200 |\n\n" + "**Totals are unaudited.**\n" + ), + "completeness": "complete", + "may_omit_content": False, + "pages_needing_ocr": [], + "truncated": False, + "parser": {"name": "anydoc", "version": "0.1.3"}, + } + assert not list(configured[0].iterdir()) + + +async def test_embedded_image_cannot_claim_complete_text(headers, configured): + image_docx = with_entry(FIXTURE.read_bytes(), "word/media/scan.png", b"\x89PNG\r\n") + response = await post( + image_docx, headers, name="report.docx", mime="application/octet-stream" + ) + assert response.status_code == 200, response.text + assert response.json()["completeness"] == "partial" + assert response.json()["may_omit_content"] is True + assert response.json()["pages_needing_ocr"] == [] + assert not list(configured[0].iterdir()) + + +def test_real_app_registers_route_only_when_enabled(): + # A fresh process verifies main.py registration, not just the test router. + script = ( + "from langchain_community.vectorstores.pgvector import PGVector\n" + "from app.services.vector_store.async_pg_vector import AsyncPgVector\n" + "PGVector.__post_init__ = lambda self: None\n" + "AsyncPgVector.__post_init__ = lambda self: None\n" + "from main import app\n" + "import sys\n" + "print(int(any(getattr(r, 'path', None) == '/v1/extract' for r in app.routes)), " + "int('anydoc' in sys.modules))\n" + ) + for enabled, expected in (("false", "0 0"), ("true", "1 0")): + env = { + **os.environ, + "RAG_EXTRACTION_API_ENABLED": enabled, + "OPENAI_API_KEY": "test_key", + } + result = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + check=True, + ) + assert result.stdout.strip().splitlines()[-1] == expected, result.stderr + + +async def test_disabled_requires_explicit_opt_in(headers, configured, monkeypatch): + monkeypatch.delenv("RAG_EXTRACTION_API_ENABLED") + response = await post(FIXTURE.read_bytes(), headers) + assert response.status_code == 404 + assert response.json()["detail"]["code"] == "EXTRACTION_DISABLED" + assert not list(configured[0].iterdir()) + + +async def test_requires_verified_identity_and_signing_secret( + headers, configured, monkeypatch +): + missing = await post(FIXTURE.read_bytes(), {}) + assert missing.status_code == 401 + monkeypatch.delenv("JWT_SECRET") + unsigned = await post(FIXTURE.read_bytes(), {}) + assert unsigned.status_code == 401 + assert unsigned.json()["detail"]["code"] == "EXTRACTION_AUTH_REQUIRED" + assert not list(configured[0].iterdir()) + + +@pytest.mark.parametrize( + "name,mime,code", + [ + ("report.md", "text/markdown", "UNSUPPORTED_DOCUMENT_TYPE"), + ("report.docx", "application/pdf", "UNSUPPORTED_DOCUMENT_TYPE"), + ("report.pdf", "application/octet-stream", "UNSUPPORTED_DOCUMENT_TYPE"), + ], +) +async def test_unrelated_formats_never_reach_the_parser( + headers, configured, name, mime, code +): + response = await post(FIXTURE.read_bytes(), headers, name=name, mime=mime) + assert response.status_code == 415 + assert response.json()["detail"]["code"] == code + assert not list(configured[0].iterdir()) + + +async def test_nonfinite_deadline_cannot_disable_worker_timeout( + headers, configured, monkeypatch +): + monkeypatch.setenv("RAG_EXTRACTION_TIMEOUT_SECONDS", "inf") + response = await post(FIXTURE.read_bytes(), headers) + assert response.status_code == 503 + assert response.json()["detail"]["code"] == "PARSER_UNAVAILABLE" + assert not list(configured[0].iterdir()) + + +async def test_unknown_profile_and_invalid_archive_fail_closed(headers, configured): + bad_profile = await post(FIXTURE.read_bytes(), headers, profile="raw-v1") + assert bad_profile.status_code == 400 + assert bad_profile.json()["detail"]["code"] == "UNSUPPORTED_PROFILE" + bad_archive = await post(b"private secret from a forged DOCX", headers) + assert bad_archive.status_code == 422 + assert bad_archive.json() == {"detail": {"code": "ARCHIVE_INVALID"}} + assert not list(configured[0].iterdir()) + + +async def test_refuses_zip_bomb_before_native_parsing(headers, configured): + bomb = with_entry( + FIXTURE.read_bytes(), "word/bomb.xml", b"x" * (25 * 1024 * 1024 + 1) + ) + response = await post(bomb, headers) + assert response.status_code == 413 + assert response.json()["detail"]["code"] == "ZIP_BOMB" + assert not list(configured[0].iterdir()) + + +async def test_empty_docx_reports_no_text_instead_of_success(headers, configured): + result = io.BytesIO() + with zipfile.ZipFile(FIXTURE) as source, zipfile.ZipFile( + result, "w", zipfile.ZIP_DEFLATED + ) as target: + for entry in source.infolist(): + data = source.read(entry) + if entry.filename == "word/document.xml": + data = ( + b'' + b"" + ) + target.writestr(entry, data) + response = await post(result.getvalue(), headers) + assert response.status_code == 422 + assert response.json()["detail"]["code"] == "NO_DOCUMENT_TEXT" + assert not list(configured[0].iterdir()) + + +async def test_archive_entry_count_has_an_independent_limit(headers, configured): + result = io.BytesIO() + with zipfile.ZipFile(FIXTURE) as source, zipfile.ZipFile( + result, "w", zipfile.ZIP_DEFLATED + ) as target: + for entry in source.infolist(): + target.writestr(entry, source.read(entry)) + for index in range(4096): + target.writestr(f"word/noise/{index}", b"") + response = await post(result.getvalue(), headers) + assert response.status_code == 413 + assert response.json()["detail"]["code"] == "ZIP_BOMB" + assert not list(configured[0].iterdir()) + + +async def test_input_limit_is_checked_while_streaming_even_without_size_hint( + configured, monkeypatch +): + monkeypatch.setattr(extraction_routes, "MAX_INPUT_BYTES", 8) + file = UploadFile(file=io.BytesIO(b"0123456789"), filename="report.docx", size=None) + with pytest.raises(HTTPException) as caught: + await extraction_routes._save_bounded(file, configured[0] / "bounded.docx") + assert caught.value.status_code == 413 + assert caught.value.detail == {"code": "PARSER_INPUT_LIMIT"} + + +async def test_worker_output_limit_is_a_refusal(configured, monkeypatch): + monkeypatch.setattr(extraction_worker, "MAX_OUTPUT_BYTES", 10) + with pytest.raises(extraction_worker.ExtractionRefusal) as caught: + extraction_worker.extract(FIXTURE) + assert caught.value.code == "PARSER_OUTPUT_LIMIT" + + +async def test_child_crash_and_malformed_output_are_sanitized( + headers, configured, monkeypatch +): + monkeypatch.setattr( + extraction, + "_command", + lambda path: (sys.executable, "-c", "import sys; sys.exit(11)"), + ) + crash = await post(FIXTURE.read_bytes(), headers) + assert crash.status_code == 503 + assert crash.json() == {"detail": {"code": "PARSER_CRASH"}} + assert not list(configured[0].iterdir()) + monkeypatch.setattr( + extraction, + "_command", + lambda path: (sys.executable, "-c", "print('private secret')"), + ) + invalid = await post(FIXTURE.read_bytes(), headers) + assert invalid.status_code == 503 + assert invalid.json() == {"detail": {"code": "PARSER_CRASH"}} + assert "private secret" not in invalid.text + assert not list(configured[0].iterdir()) + + +async def test_api_kills_worker_that_overproduces_ipc(headers, configured, monkeypatch): + monkeypatch.setattr(extraction, "MAX_IPC_BYTES", 64) + monkeypatch.setattr( + extraction, + "_command", + lambda path: (sys.executable, "-c", "import os; os.write(1, b'x' * 1024)"), + ) + response = await post(FIXTURE.read_bytes(), headers) + assert response.status_code == 413 + assert response.json()["detail"]["code"] == "PARSER_OUTPUT_LIMIT" + assert not list(configured[0].iterdir()) + assert configured[1]._pending == 0 + + +def _sleeping_command(path: Path) -> tuple[str, ...]: + return ( + sys.executable, + "-c", + "import os,sys,time; open(sys.argv[1]+'.pid','w').write(str(os.getpid())); time.sleep(60)", + str(path), + ) + + +async def _await_child(tmp_path: Path) -> Path: + for _ in range(200): + pids = list(tmp_path.glob("*.pid")) + if pids: + return pids[0] + await asyncio.sleep(0.01) + raise AssertionError("parser child did not start") + + +async def _assert_child_reaped( + pid_file: Path, admission: extraction.ExtractionAdmission +) -> None: + pid = int(pid_file.read_text()) + for _ in range(200): + try: + os.kill(pid, 0) + except ProcessLookupError: + if admission._pending == 0 and not list(pid_file.parent.glob("*.docx")): + pid_file.unlink() + return + await asyncio.sleep(0.01) + raise AssertionError("parser child or staging file survived cancellation") + + +async def test_timeout_kills_and_reaps_child(headers, configured, monkeypatch): + monkeypatch.setattr(extraction, "_command", _sleeping_command) + monkeypatch.setenv("RAG_EXTRACTION_TIMEOUT_SECONDS", "0.3") + task = asyncio.create_task(post(FIXTURE.read_bytes(), headers)) + pid_file = await _await_child(configured[0]) + response = await asyncio.wait_for(task, 2) + assert response.status_code == 504 + assert response.json()["detail"]["code"] == "PARSER_TIMEOUT" + await _assert_child_reaped(pid_file, configured[1]) + assert not list(configured[0].iterdir()) + assert configured[1]._pending == 0 + + +async def test_cancel_reaps_child_and_frees_slot_for_next_upload( + headers, configured, monkeypatch +): + original_command = extraction._command + monkeypatch.setattr(extraction, "_command", _sleeping_command) + task = asyncio.create_task(post(FIXTURE.read_bytes(), headers)) + pid_file = await _await_child(configured[0]) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + await _assert_child_reaped(pid_file, configured[1]) + assert not list(configured[0].iterdir()) + assert configured[1]._pending == 0 + monkeypatch.setattr(extraction, "_command", original_command) + response = await post(FIXTURE.read_bytes(), headers) + assert response.status_code == 200 + + +async def test_busy_parser_refuses_before_staging_anything( + headers, configured, monkeypatch +): + monkeypatch.setattr(extraction, "_command", _sleeping_command) + monkeypatch.setattr( + extraction_routes, "_admission", extraction.ExtractionAdmission(1, 0) + ) + task = asyncio.create_task(post(FIXTURE.read_bytes(), headers)) + pid_file = await _await_child(configured[0]) + busy = await post(FIXTURE.read_bytes(), headers) + assert busy.status_code == 429 + assert busy.json()["detail"]["code"] == "CONCURRENCY_LIMIT" + assert len(list(configured[0].glob("*.docx"))) == 1 + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 2) + await _assert_child_reaped(pid_file, extraction_routes._admission) + assert not list(configured[0].iterdir()) + + +async def test_existing_text_endpoint_still_preserves_raw_markdown(headers, configured): + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=main_app), base_url="http://test" + ) as client: + response = await client.post( + "/text", + headers=headers, + data={"file_id": "md-test"}, + files={"file": ("readme.md", b"# Raw **Markdown**\n", "text/markdown")}, + ) + assert response.status_code == 200, response.text + assert response.json()["text"] == "# Raw **Markdown**"