diff --git a/.env.example b/.env.example index 312db5114..71f64d116 100644 --- a/.env.example +++ b/.env.example @@ -84,6 +84,18 @@ API_BASE_URL='http://X.X.X.X:APP_PORT' # Base URL of your FastAPI backe # WEBSEARCH_TOP_K=5 # Number of web results to include (default: 5) # WEBSEARCH_LANG=fr-FR # Search language/market (default: fr-FR) +# MCP SERVER +# Standalone Model Context Protocol server (openrag/api/mcp/server.py). +# OPENRAG_MCP_SERVER_NAME="OpenRAG MCP" +# OPENRAG_MCP_HOST=0.0.0.0 +# OPENRAG_MCP_PORT=8081 +# OPENRAG_MCP_PATH=/mcp +# OPENRAG_MCP_DEFAULT_TOP_K=5 +# OPENRAG_MCP_MAX_TOP_K=50 +# OPENRAG_MCP_SIMILARITY_THRESHOLD=0.8 +# OPENRAG_MCP_DOWNLOAD_TIMEOUT=30.0 +# OPENRAG_MCP_MAX_DOWNLOAD_BYTES=104857600 # 100 MiB + # LOGGING LOG_LEVEL=DEBUG # See possible values https://loguru.readthedocs.io/en/stable/api/logger.html diff --git a/conf/config.yaml b/conf/config.yaml index 393c0dfab..4f058accb 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -320,3 +320,19 @@ websearch: fetch_timeout: 1.0 fetch_max_tokens: 500 fetch_verify_ssl: false + +# --- MCP server --- +# Env: OPENRAG_MCP_SERVER_NAME, OPENRAG_MCP_HOST, OPENRAG_MCP_PORT, +# OPENRAG_MCP_PATH, OPENRAG_MCP_DEFAULT_TOP_K, OPENRAG_MCP_MAX_TOP_K, +# OPENRAG_MCP_SIMILARITY_THRESHOLD, OPENRAG_MCP_DOWNLOAD_TIMEOUT, +# OPENRAG_MCP_MAX_DOWNLOAD_BYTES +mcp: + server_name: OpenRAG MCP + host: 0.0.0.0 + port: 8081 + path: /mcp + default_top_k: 5 + max_top_k: 50 + similarity_threshold: 0.8 + download_timeout: 30.0 + max_download_bytes: 104857600 # 100 MiB diff --git a/openrag/api/mcp/__init__.py b/openrag/api/mcp/__init__.py new file mode 100644 index 000000000..85cd5442c --- /dev/null +++ b/openrag/api/mcp/__init__.py @@ -0,0 +1 @@ +"""Standalone MCP (Model Context Protocol) server for OpenRAG.""" diff --git a/openrag/api/mcp/auth_context.py b/openrag/api/mcp/auth_context.py new file mode 100644 index 000000000..97c3a7f7b --- /dev/null +++ b/openrag/api/mcp/auth_context.py @@ -0,0 +1,51 @@ +"""Request-scoped auth context for the MCP server. + +FastMCP tool callables receive only their declared arguments — there is no +FastAPI ``Request`` to thread the authenticated principal through. The +``MCPAuthContextMiddleware`` resolves the caller once per request and stashes +the principal here; the tool functions read it and forward it explicitly to +``MCPService`` (which never touches this module — it stays a pure +orchestrator). +""" + +from __future__ import annotations + +from contextvars import ContextVar, Token + +_USER_ID: ContextVar[int | None] = ContextVar("openrag_mcp_user_id", default=None) +_IS_ADMIN: ContextVar[bool] = ContextVar("openrag_mcp_is_admin", default=False) +_ALLOWED_PARTITIONS: ContextVar[list[str] | None] = ContextVar("openrag_mcp_partitions", default=None) + +AuthTokens = tuple[Token, Token, Token] + + +def set_auth_context( + *, + user_id: int | None, + is_admin: bool, + allowed_partitions: list[str] | None, +) -> AuthTokens: + return ( + _USER_ID.set(user_id), + _IS_ADMIN.set(is_admin), + _ALLOWED_PARTITIONS.set(allowed_partitions), + ) + + +def reset_auth_context(tokens: AuthTokens) -> None: + user_token, admin_token, partitions_token = tokens + _USER_ID.reset(user_token) + _IS_ADMIN.reset(admin_token) + _ALLOWED_PARTITIONS.reset(partitions_token) + + +def get_user_id() -> int | None: + return _USER_ID.get() + + +def is_admin() -> bool: + return _IS_ADMIN.get() + + +def get_allowed_partitions() -> list[str] | None: + return _ALLOWED_PARTITIONS.get() diff --git a/openrag/api/mcp/server.py b/openrag/api/mcp/server.py new file mode 100644 index 000000000..5b924eb02 --- /dev/null +++ b/openrag/api/mcp/server.py @@ -0,0 +1,471 @@ +"""Standalone MCP (Model Context Protocol) server. + +Reimplementation of PR 273's ``openrag/mcp_server.py`` on the hexagonal +stack. A FastMCP streamable-HTTP app exposes OpenRAG's search / catalog / +indexation operations as MCP tools. All tool logic lives in +``services.orchestrators.mcp_service.MCPService``; this module is the thin +``api`` entrypoint that: + +* owns a :class:`~di.container.ServiceContainer` for the server process + (built and initialised in the ASGI lifespan, mirroring ``api/main.py``); +* resolves the caller from the bearer token in + :class:`MCPAuthContextMiddleware` and stashes the principal in + :mod:`api.mcp.auth_context`; +* wires one ``@server.tool`` per operation, each reading the auth context + and forwarding it explicitly to ``MCPService``. + +Per the layer guard, ``api`` reaches ``services`` only through ``di`` — the +``MCPService`` instance is pulled off the container (``container.mcp_service``) +and never imported here. + +Run it with ``uv run -m api.mcp.server`` (serves at ``http://host:port``, +defaults ``0.0.0.0:8081/mcp``; configure via the ``OPENRAG_MCP_*`` env vars). +""" + +from __future__ import annotations + +import os +from contextlib import asynccontextmanager + +import ray +from api.mcp.auth_context import ( + get_allowed_partitions, + get_user_id, + is_admin, + reset_auth_context, + set_auth_context, +) +from config import load_config +from core.utils.log_tail import app_log_file +from di.container import ServiceContainer +from di.workers import ensure_worker_bootstrap +from mcp.server.fastmcp import FastMCP +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse +from utils.logger import get_logger + +logger = get_logger() + +config = load_config() +mcp_config = config.mcp + +server = FastMCP(mcp_config.server_name, stateless_http=True, json_response=True) + +LOG_FILE = app_log_file(config.paths.log_dir) + + +# --------------------------------------------------------------------------- +# Process-level container (built in the ASGI lifespan, on the serving loop) +# --------------------------------------------------------------------------- + +_container: ServiceContainer | None = None + + +def _require_container() -> ServiceContainer: + if _container is None: + raise RuntimeError("MCP service container is not initialized.") + return _container + + +def _service(): + """The :class:`MCPService` for the current process (typed loosely so the + ``api`` layer does not import from ``services``).""" + return _require_container().mcp_service + + +async def _startup() -> None: + global _container + if not ray.is_initialized(): + ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True) + ensure_worker_bootstrap() + container = ServiceContainer(config) + await container.initialize() + _container = container + logger.info("MCP server container initialized") + + +async def _shutdown() -> None: + global _container + if _container is not None: + try: + await _container.shutdown() + finally: + _container = None + + +# --------------------------------------------------------------------------- +# Auth middleware — resolves the caller and stashes the principal +# --------------------------------------------------------------------------- + + +class MCPAuthContextMiddleware(BaseHTTPMiddleware): + """Resolve the bearer-token principal and set the MCP auth context. + + Mirrors the token-mode contract of ``api.middleware.AuthMiddleware``: + when ``AUTH_TOKEN`` is unset the server runs in dev mode and every call + acts as admin user 1; otherwise a valid ``Authorization: Bearer`` is + required. Allowed partitions follow ``current_user_or_admin_partitions_list`` + (``["all"]`` for admins under ``SUPER_ADMIN_MODE``). + """ + + async def _resolve_principal(self, request: Request) -> tuple[int | None, bool, list[str] | None]: + auth_service = _require_container().auth_service + auth_mode = os.getenv("AUTH_MODE", "token").strip().lower() + auth_token = os.getenv("AUTH_TOKEN") + + # Dev bypass (act as admin user 1) is allowed ONLY in token mode with no + # configured admin token — identical to api.middleware.AuthMiddleware. + # In OIDC mode AUTH_TOKEN is commonly unset; without this gate the MCP + # endpoint would silently treat every caller as admin (auth bypass), so + # a valid bearer (users.token) is always required there. + if auth_mode == "token" and auth_token is None: + user = await auth_service.get_user_for_request(1) + else: + header = request.headers.get("authorization", "") + token = header.split(" ", 1)[1] if header.lower().startswith("bearer ") else None + if not token: + raise PermissionError("Missing token") + user = await auth_service.get_user_by_token_for_request(token) + if not user: + raise PermissionError("Invalid token") + if not user: + raise PermissionError("Unauthenticated") + + super_admin_mode = os.getenv("SUPER_ADMIN_MODE", "false").lower() == "true" + admin = bool(user.get("is_admin")) + if admin and super_admin_mode: + allowed = ["all"] + else: + user_partitions = await auth_service.list_user_partitions_for_request(user["id"]) + allowed = [p["partition"] for p in user_partitions] + return user.get("id"), admin, allowed + + async def dispatch(self, request: Request, call_next): + try: + user_id, admin, allowed = await self._resolve_principal(request) + except PermissionError as exc: + return JSONResponse(status_code=403, content={"detail": str(exc)}) + except Exception: # pragma: no cover - defensive + # Don't leak internal exception detail to the client. + logger.exception("MCP authentication error") + return JSONResponse(status_code=500, content={"detail": "Internal authentication error"}) + + tokens = set_auth_context(user_id=user_id, is_admin=admin, allowed_partitions=allowed) + try: + return await call_next(request) + finally: + reset_auth_context(tokens) + + +# --------------------------------------------------------------------------- +# Tools — search +# --------------------------------------------------------------------------- + + +@server.tool(description="Semantic search across one or many partitions using OpenRAG search flow") +async def search_documents(query: str, partitions: list[str] | None = None, top_k: int | None = None) -> dict: + return await _service().search_documents( + query=query, + partitions=partitions, + top_k=top_k, + allowed_partitions=get_allowed_partitions(), + ) + + +@server.tool(description="Semantic search restricted to one partition") +async def search_partition(query: str, partition: str, top_k: int | None = None) -> dict: + return await _service().search_documents( + query=query, + partitions=[partition], + top_k=top_k, + allowed_partitions=get_allowed_partitions(), + ) + + +@server.tool(description="Semantic search restricted to one file inside one partition") +async def search_file(query: str, partition: str, file_id: str, top_k: int | None = None) -> dict: + return await _service().search_documents( + query=query, + partitions=[partition], + top_k=top_k, + file_id=file_id, + allowed_partitions=get_allowed_partitions(), + ) + + +# --------------------------------------------------------------------------- +# Tools — catalog browsing +# --------------------------------------------------------------------------- + + +@server.tool( + description=( + "List all partitions accessible to the current user. " + "Returns partition names, creation timestamps, and total count." + ) +) +async def list_partitions() -> dict: + return await _service().list_partitions(allowed_partitions=get_allowed_partitions()) + + +@server.tool( + description=( + "List all files indexed in a given partition. " + "Returns file IDs, original filenames, sizes, creation dates, and other metadata. " + "Use `limit` to cap the number of results." + ) +) +async def list_files(partition: str, limit: int | None = None) -> dict: + return await _service().list_files( + partition=partition, + allowed_partitions=get_allowed_partitions(), + limit=limit, + ) + + +@server.tool( + description=( + "Get metadata and chunk count for a specific file inside a partition. " + "Returns file metadata (filename, size, creation date, …) and the total number of indexed chunks." + ) +) +async def get_file_info(partition: str, file_id: str) -> dict: + return await _service().get_file_info( + partition=partition, + file_id=file_id, + allowed_partitions=get_allowed_partitions(), + ) + + +@server.tool( + description=( + "Fetch text chunks belonging to a specific file, one page at a time. " + "Use `offset` (default 0) and `limit` (default 3) to page through the file. " + "Keep `limit` small (3 or fewer) to avoid exceeding your context window — " + "chunks can be long. " + "The response includes `total_chunks` and `has_more` so you know whether to " + "call again with a higher offset. Always check `has_more` and keep paging " + "until it is false before drawing conclusions about the full file content." + ) +) +async def get_file_chunks(partition: str, file_id: str, offset: int = 0, limit: int = 3) -> dict: + return await _service().get_file_chunks( + partition=partition, + file_id=file_id, + allowed_partitions=get_allowed_partitions(), + offset=offset, + limit=limit, + ) + + +@server.tool( + description=( + "Fuzzy search across file names (filename, original_filename, file_id) " + "using sequence-similarity matching. " + "Results are ranked by similarity score (0–1). " + "Optionally restrict the search to a single `partition`. " + "Use `cutoff` (default 0.4) to control minimum similarity and `limit` (default 20) to cap results." + ) +) +async def fuzzy_search_files(query: str, partition: str | None = None, cutoff: float = 0.4, limit: int = 20) -> dict: + return await _service().fuzzy_search_files( + query=query, + allowed_partitions=get_allowed_partitions(), + partition=partition, + cutoff=cutoff, + limit=limit, + ) + + +@server.tool( + description=( + "Fetch a single indexed chunk by its ID. " + "Chunk IDs appear in search results and in get_file_chunks output. " + "Returns the full text content and all metadata for that chunk." + ) +) +async def get_chunk_by_id(chunk_id: str) -> dict: + return await _service().get_chunk_by_id( + chunk_id=chunk_id, + allowed_partitions=get_allowed_partitions(), + ) + + +# --------------------------------------------------------------------------- +# Tools — task introspection +# --------------------------------------------------------------------------- + + +@server.tool( + description=( + "Get the current status and details of an indexation task. " + "Task states: QUEUED → SERIALIZING → CHUNKING → INSERTING → COMPLETED (or FAILED). " + "If the task failed, the error message is included in the response." + ) +) +async def get_indexation_task_status(task_id: str) -> dict: + return await _service().get_task_status( + task_id=task_id, + user_id=get_user_id(), + is_admin=is_admin(), + ) + + +@server.tool( + description=( + "List all indexation tasks belonging to the current user. " + "Use `task_status` to filter: 'active' (queued/in-progress), 'completed', 'failed', " + "or any exact state name. Omit to get all tasks." + ) +) +async def list_my_tasks(task_status: str | None = None) -> dict: + return await _service().list_my_tasks( + user_id=get_user_id(), + is_admin=is_admin(), + task_status=task_status, + ) + + +@server.tool( + description=( + "Fetch chronological log lines for a specific indexation task. " + "Useful for diagnosing slow or stuck indexations. " + "Use `max_lines` to cap output (default 100)." + ) +) +async def get_task_logs(task_id: str, max_lines: int = 100) -> dict: + return await _service().get_task_logs( + task_id=task_id, + user_id=get_user_id(), + is_admin=is_admin(), + log_file=LOG_FILE, + max_lines=max_lines, + ) + + +# --------------------------------------------------------------------------- +# Tools — write operations +# --------------------------------------------------------------------------- + + +@server.tool( + description=( + "Delete a file and all its indexed chunks from a partition. " + "Requires editor (or owner) access to the partition. " + "This operation is irreversible." + ) +) +async def delete_file(partition: str, file_id: str) -> dict: + return await _service().delete_file( + partition=partition, + file_id=file_id, + allowed_partitions=get_allowed_partitions(), + user_id=get_user_id(), + ) + + +@server.tool( + description=( + "Update metadata fields of an existing indexed file without re-uploading it. " + "Pass a JSON object with only the fields to change (e.g. author, title). " + "To move the file to another partition, include a 'partition' key — " + "you must have editor access to both the source and destination partitions." + ) +) +async def update_file_metadata(partition: str, file_id: str, metadata: dict) -> dict: + return await _service().update_file_metadata( + partition=partition, + file_id=file_id, + metadata=metadata, + allowed_partitions=get_allowed_partitions(), + user_id=get_user_id(), + ) + + +@server.tool( + description=( + "Copy a file from one partition to another. " + "Requires read access to the source partition and editor access to the destination. " + "Optionally supply `extra_metadata` to override fields in the copy." + ) +) +async def copy_file( + source_partition: str, + source_file_id: str, + dest_partition: str, + dest_file_id: str, + extra_metadata: dict | None = None, +) -> dict: + return await _service().copy_file( + source_partition=source_partition, + source_file_id=source_file_id, + dest_partition=dest_partition, + dest_file_id=dest_file_id, + allowed_partitions=get_allowed_partitions(), + user_id=get_user_id(), + extra_metadata=extra_metadata, + ) + + +@server.tool( + description=( + "Download a document from a public HTTP/HTTPS URL and index it into a partition. " + "Returns a task_id that can be polled with get_indexation_task_status. " + "The file_id must be unique within the partition. " + "Optionally supply `extra_metadata` (dict) to attach custom fields." + ) +) +async def index_url(url: str, partition: str, file_id: str, extra_metadata: dict | None = None) -> dict: + return await _service().index_url( + url=url, + partition=partition, + file_id=file_id, + allowed_partitions=get_allowed_partitions(), + user_id=get_user_id(), + extra_metadata=extra_metadata, + ) + + +# --------------------------------------------------------------------------- +# App factory + entrypoint +# --------------------------------------------------------------------------- + + +def create_mcp_http_app(): + """Build the FastMCP streamable-HTTP app with auth + container lifecycle. + + The container is initialised on the serving event loop (asyncpg pools are + loop-bound) by composing our startup/shutdown around FastMCP's own + session-manager lifespan. + """ + server.settings.host = mcp_config.host + server.settings.port = mcp_config.port + server.settings.streamable_http_path = mcp_config.path + + app = server.streamable_http_app() + app.add_middleware(MCPAuthContextMiddleware) + + fastmcp_lifespan = app.router.lifespan_context + + @asynccontextmanager + async def _combined_lifespan(app_): + await _startup() + try: + async with fastmcp_lifespan(app_): + yield + finally: + await _shutdown() + + app.router.lifespan_context = _combined_lifespan + return app + + +def main(): + import uvicorn + + uvicorn.run(create_mcp_http_app(), host=mcp_config.host, port=mcp_config.port) + + +if __name__ == "__main__": + main() diff --git a/openrag/api/mcp/test_server.py b/openrag/api/mcp/test_server.py new file mode 100644 index 000000000..07939df4e --- /dev/null +++ b/openrag/api/mcp/test_server.py @@ -0,0 +1,240 @@ +"""Tests for the MCP server entrypoint: auth context, middleware, tool wiring.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from api.mcp import auth_context as ac +from api.mcp import server +from starlette.requests import Request +from starlette.responses import Response + +# --------------------------------------------------------------------------- +# auth_context +# --------------------------------------------------------------------------- + + +def test_auth_context_roundtrip(): + assert ac.get_user_id() is None + assert ac.is_admin() is False + assert ac.get_allowed_partitions() is None + + tokens = ac.set_auth_context(user_id=7, is_admin=True, allowed_partitions=["a"]) + assert ac.get_user_id() == 7 + assert ac.is_admin() is True + assert ac.get_allowed_partitions() == ["a"] + + ac.reset_auth_context(tokens) + assert ac.get_user_id() is None + assert ac.is_admin() is False + assert ac.get_allowed_partitions() is None + + +# --------------------------------------------------------------------------- +# Middleware fakes + helpers +# --------------------------------------------------------------------------- + + +class FakeAuthService: + def __init__(self, *, user=None, by_token=None, partitions=None): + self._user = user + self._by_token = by_token + self._partitions = partitions if partitions is not None else [] + + async def get_user_for_request(self, user_id): + return self._user + + async def get_user_by_token_for_request(self, token): + return self._by_token + + async def list_user_partitions_for_request(self, user_id): + return list(self._partitions) + + +def _install_container(monkeypatch, auth_service): + monkeypatch.setattr(server, "_container", SimpleNamespace(auth_service=auth_service)) + + +def _request(headers=None): + raw = [(k.lower().encode(), v.encode()) for k, v in (headers or {}).items()] + scope = {"type": "http", "method": "GET", "path": "/mcp", "headers": raw, "query_string": b""} + return Request(scope) + + +async def _dispatch(monkeypatch, request): + """Run the middleware and return (response, captured_context).""" + captured: dict = {} + + async def call_next(_req): + captured["user_id"] = ac.get_user_id() + captured["is_admin"] = ac.is_admin() + captured["allowed"] = ac.get_allowed_partitions() + return Response("ok") + + mw = server.MCPAuthContextMiddleware(lambda scope, receive, send: None) + response = await mw.dispatch(request, call_next) + # context must be reset after dispatch regardless of outcome + assert ac.get_user_id() is None + assert ac.get_allowed_partitions() is None + return response, captured + + +# --------------------------------------------------------------------------- +# Middleware behaviour +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dev_mode_resolves_user_one(monkeypatch): + monkeypatch.setenv("AUTH_MODE", "token") + monkeypatch.delenv("AUTH_TOKEN", raising=False) + monkeypatch.delenv("SUPER_ADMIN_MODE", raising=False) + _install_container( + monkeypatch, + FakeAuthService(user={"id": 1, "is_admin": True}, partitions=[{"partition": "a"}, {"partition": "b"}]), + ) + + response, captured = await _dispatch(monkeypatch, _request()) + + assert response.status_code == 200 + assert captured["user_id"] == 1 + # admin but SUPER_ADMIN_MODE off → explicit partition list, not wildcard + assert captured["allowed"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_super_admin_gets_wildcard(monkeypatch): + monkeypatch.setenv("AUTH_MODE", "token") + monkeypatch.delenv("AUTH_TOKEN", raising=False) + monkeypatch.setenv("SUPER_ADMIN_MODE", "true") + _install_container(monkeypatch, FakeAuthService(user={"id": 1, "is_admin": True}, partitions=[{"partition": "a"}])) + + _, captured = await _dispatch(monkeypatch, _request()) + + assert captured["is_admin"] is True + assert captured["allowed"] == ["all"] + + +@pytest.mark.asyncio +async def test_token_mode_missing_token_403(monkeypatch): + monkeypatch.setenv("AUTH_TOKEN", "secret") + _install_container(monkeypatch, FakeAuthService()) + + response, captured = await _dispatch(monkeypatch, _request()) + + assert response.status_code == 403 + assert captured == {} # call_next never ran + + +@pytest.mark.asyncio +async def test_token_mode_invalid_token_403(monkeypatch): + monkeypatch.setenv("AUTH_TOKEN", "secret") + _install_container(monkeypatch, FakeAuthService(by_token=None)) + + response, _ = await _dispatch(monkeypatch, _request(headers={"authorization": "Bearer nope"})) + + assert response.status_code == 403 + + +@pytest.mark.asyncio +async def test_oidc_mode_no_token_is_not_dev_bypassed(monkeypatch): + # H1 regression: AUTH_MODE=oidc with AUTH_TOKEN unset must NOT fall into the + # user-1 admin dev bypass — a missing bearer is rejected. + monkeypatch.setenv("AUTH_MODE", "oidc") + monkeypatch.delenv("AUTH_TOKEN", raising=False) + _install_container(monkeypatch, FakeAuthService(user={"id": 1, "is_admin": True})) + + response, captured = await _dispatch(monkeypatch, _request()) + + assert response.status_code == 403 + assert captured == {} # never reached the tool as admin + + +@pytest.mark.asyncio +async def test_oidc_mode_valid_bearer_sets_context(monkeypatch): + monkeypatch.setenv("AUTH_MODE", "oidc") + monkeypatch.delenv("AUTH_TOKEN", raising=False) + monkeypatch.delenv("SUPER_ADMIN_MODE", raising=False) + _install_container( + monkeypatch, + FakeAuthService(by_token={"id": 8, "is_admin": False}, partitions=[{"partition": "p"}]), + ) + + response, captured = await _dispatch(monkeypatch, _request(headers={"authorization": "Bearer good"})) + + assert response.status_code == 200 + assert captured["user_id"] == 8 + assert captured["allowed"] == ["p"] + + +@pytest.mark.asyncio +async def test_token_mode_valid_token_sets_context(monkeypatch): + monkeypatch.setenv("AUTH_TOKEN", "secret") + monkeypatch.delenv("SUPER_ADMIN_MODE", raising=False) + _install_container( + monkeypatch, + FakeAuthService(by_token={"id": 42, "is_admin": False}, partitions=[{"partition": "team"}]), + ) + + response, captured = await _dispatch(monkeypatch, _request(headers={"authorization": "Bearer good"})) + + assert response.status_code == 200 + assert captured["user_id"] == 42 + assert captured["is_admin"] is False + assert captured["allowed"] == ["team"] + + +# --------------------------------------------------------------------------- +# Tool → service wiring +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_search_tool_forwards_auth_context(monkeypatch): + calls: dict = {} + + async def search_documents(**kwargs): + calls.update(kwargs) + return {"ok": True} + + fake_mcp = SimpleNamespace(search_documents=search_documents) + monkeypatch.setattr(server, "_container", SimpleNamespace(mcp_service=fake_mcp)) + + tokens = ac.set_auth_context(user_id=3, is_admin=False, allowed_partitions=["a"]) + try: + out = await server.search_documents(query="hi", top_k=4) + finally: + ac.reset_auth_context(tokens) + + assert out == {"ok": True} + assert calls["query"] == "hi" + assert calls["top_k"] == 4 + assert calls["allowed_partitions"] == ["a"] + + +@pytest.mark.asyncio +async def test_delete_tool_forwards_user_id(monkeypatch): + calls: dict = {} + + async def delete_file(**kwargs): + calls.update(kwargs) + return {"ok": True} + + monkeypatch.setattr(server, "_container", SimpleNamespace(mcp_service=SimpleNamespace(delete_file=delete_file))) + + tokens = ac.set_auth_context(user_id=9, is_admin=False, allowed_partitions=["p"]) + try: + await server.delete_file(partition="p", file_id="f") + finally: + ac.reset_auth_context(tokens) + + assert calls["user_id"] == 9 + assert calls["partition"] == "p" + assert calls["allowed_partitions"] == ["p"] + + +def test_require_container_raises_when_unset(monkeypatch): + monkeypatch.setattr(server, "_container", None) + with pytest.raises(RuntimeError): + server._require_container() diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index 3a70b9839..188ccd43e 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -29,6 +29,7 @@ ) from api.routers.admin.task_logs import collect_task_logs from components.indexer.utils.files import sanitize_filename, save_file_to_disk +from core.utils.log_tail import app_log_file from di.providers import get_auth_service, get_config, get_indexing_service, get_partition_service from fastapi import ( APIRouter, @@ -492,7 +493,7 @@ async def get_task_logs( task_details=Depends(require_task_owner), config=Depends(get_config), ): - log_file = Path(config.paths.log_dir or "logs") / "app.json" + log_file = app_log_file(config.paths.log_dir) if not log_file.exists(): raise HTTPException(status_code=500, detail="Log file not found.") diff --git a/openrag/api/routers/admin/task_logs.py b/openrag/api/routers/admin/task_logs.py index c7e854e3b..d6356cc1d 100644 --- a/openrag/api/routers/admin/task_logs.py +++ b/openrag/api/routers/admin/task_logs.py @@ -1,52 +1,10 @@ -import json -from pathlib import Path - -MAX_TASK_LOG_LINES = 5000 -TASK_LOG_READ_BLOCK_SIZE = 64 * 1024 - - -def iter_file_lines_reversed(path: Path, block_size: int = TASK_LOG_READ_BLOCK_SIZE): - with path.open("rb") as f: - f.seek(0, 2) - position = f.tell() - pending = b"" - - while position > 0: - read_size = min(block_size, position) - position -= read_size - f.seek(position) - pending = f.read(read_size) + pending - lines = pending.split(b"\n") - pending = lines[0] - - for line in reversed(lines[1:]): - if line: - yield line.decode(errors="replace") - - if pending: - yield pending.decode(errors="replace") - - -def collect_task_logs(log_file: Path, task_id: str, max_lines: int) -> list[str]: - if max_lines < 1 or max_lines > MAX_TASK_LOG_LINES: - raise ValueError(f"max_lines must be between 1 and {MAX_TASK_LOG_LINES}") - - logs = [] - for line in iter_file_lines_reversed(log_file): - try: - record = json.loads(line).get("record", {}) - extra = record.get("extra") or {} - if extra.get("task_id") != task_id: - continue - time_repr = (record.get("time") or {}).get("repr") - level_name = (record.get("level") or {}).get("name") - message = record.get("message") - if not all((time_repr, level_name, message)): - continue - logs.append(f"{time_repr} - {level_name} - {message} - {extra}") - if len(logs) >= max_lines: - break - except (json.JSONDecodeError, AttributeError): - continue - - return logs[::-1] +# Re-exported for callers/tests that import these from this module; the single +# implementation now lives in core.utils so the api route and the MCP server +# share it. +from core.utils.log_tail import ( + MAX_TASK_LOG_LINES, + collect_task_logs, + iter_file_lines_reversed, +) + +__all__ = ["MAX_TASK_LOG_LINES", "collect_task_logs", "iter_file_lines_reversed"] diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py index 030719519..bb7dccc5a 100644 --- a/openrag/core/config/loader.py +++ b/openrag/core/config/loader.py @@ -169,6 +169,16 @@ ("WEBSEARCH_FETCH_TIMEOUT", "websearch.fetch_timeout", float), ("WEBSEARCH_FETCH_MAX_TOKENS", "websearch.fetch_max_tokens", int), ("WEBSEARCH_FETCH_VERIFY_SSL", "websearch.fetch_verify_ssl", bool), + # MCP server + ("OPENRAG_MCP_SERVER_NAME", "mcp.server_name", str), + ("OPENRAG_MCP_HOST", "mcp.host", str), + ("OPENRAG_MCP_PORT", "mcp.port", int), + ("OPENRAG_MCP_PATH", "mcp.path", str), + ("OPENRAG_MCP_DEFAULT_TOP_K", "mcp.default_top_k", int), + ("OPENRAG_MCP_MAX_TOP_K", "mcp.max_top_k", int), + ("OPENRAG_MCP_SIMILARITY_THRESHOLD", "mcp.similarity_threshold", float), + ("OPENRAG_MCP_DOWNLOAD_TIMEOUT", "mcp.download_timeout", float), + ("OPENRAG_MCP_MAX_DOWNLOAD_BYTES", "mcp.max_download_bytes", int), ] _AUDIO_EXTENSIONS = ("mp3", "flac", "ogg", "aac", "flv", "wma", "mp4") diff --git a/openrag/core/config/mcp.py b/openrag/core/config/mcp.py new file mode 100644 index 000000000..8d9af7fc4 --- /dev/null +++ b/openrag/core/config/mcp.py @@ -0,0 +1,24 @@ +"""MCP server configuration. + +Settings for the standalone Model Context Protocol server +(``openrag/api/mcp/server.py``): the FastMCP transport binding plus the +search-tool defaults/bounds applied before a request reaches +``RetrievalService``. +""" + +from __future__ import annotations + +from .base import ConfigMixin + + +class MCPServerConfig(ConfigMixin): + server_name: str = "OpenRAG MCP" + host: str = "0.0.0.0" + port: int = 8081 + path: str = "/mcp" + default_top_k: int = 5 + max_top_k: int = 50 + similarity_threshold: float = 0.8 + # Bounds for the index_url server-side fetch (SSRF/DoS hardening). + download_timeout: float = 30.0 + max_download_bytes: int = 100 * 1024 * 1024 # 100 MiB diff --git a/openrag/core/config/root.py b/openrag/core/config/root.py index cbb219a43..c44483ace 100644 --- a/openrag/core/config/root.py +++ b/openrag/core/config/root.py @@ -23,6 +23,7 @@ VectorDBConfig, VerboseConfig, ) +from .mcp import MCPServerConfig from .retrieval import ( MapReduceConfig, RAGConfig, @@ -61,3 +62,4 @@ class Settings(ConfigMixin): retriever: RetrieverConfig = Field(default_factory=SingleRetrieverConfig) rag: RAGConfig = Field(default_factory=RAGConfig) websearch: WebSearchConfig = Field(default_factory=StaanWebSearchConfig) + mcp: MCPServerConfig = Field(default_factory=MCPServerConfig) diff --git a/openrag/core/utils/log_tail.py b/openrag/core/utils/log_tail.py new file mode 100644 index 000000000..838a5318a --- /dev/null +++ b/openrag/core/utils/log_tail.py @@ -0,0 +1,92 @@ +"""Locate and tail the application log. + +The single source of truth for the JSON application-log path (``app.json`` +under the configured ``log_dir``), the backwards block reader that collects +the newest matching lines without loading the whole file, and the loguru-record +parser that surfaces one task's lines. Shared by the logger sink, the admin +task-logs route (``api``) and the MCP ``get_task_logs`` tool (``services``) — +the single implementation lives here so no layer duplicates it. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from pathlib import Path + +DEFAULT_BLOCK_SIZE = 64 * 1024 +MAX_TASK_LOG_LINES = 5000 +APP_LOG_FILENAME = "app.json" + + +def app_log_file(log_dir: str | Path | None) -> Path: + """Canonical path of the JSON application log under *log_dir*. + + Falls back to ``logs/`` when *log_dir* is unset/empty. + """ + return Path(log_dir or "logs") / APP_LOG_FILENAME + + +def iter_file_lines_reversed(path: Path, block_size: int = DEFAULT_BLOCK_SIZE) -> Iterator[str]: + """Yield the file's lines newest-first, reading from the end in blocks.""" + with path.open("rb") as f: + f.seek(0, 2) + position = f.tell() + pending = b"" + + while position > 0: + read_size = min(block_size, position) + position -= read_size + f.seek(position) + pending = f.read(read_size) + pending + lines = pending.split(b"\n") + pending = lines[0] + + for line in reversed(lines[1:]): + if line: + yield line.decode(errors="replace") + + if pending: + yield pending.decode(errors="replace") + + +def collect_task_logs(log_file: Path, task_id: str, max_lines: int) -> list[str]: + """Return chronological loguru log lines tagged with ``task_id``. + + Scans the JSON-per-line application log newest-first (block by block, never + loading the whole file) until ``max_lines`` matching records are collected, + then restores chronological order. ``max_lines`` is bounded by + ``MAX_TASK_LOG_LINES``. + """ + if max_lines < 1 or max_lines > MAX_TASK_LOG_LINES: + raise ValueError(f"max_lines must be between 1 and {MAX_TASK_LOG_LINES}") + + logs: list[str] = [] + for line in iter_file_lines_reversed(log_file): + try: + record = json.loads(line).get("record", {}) + extra = record.get("extra") or {} + if extra.get("task_id") != task_id: + continue + time_repr = (record.get("time") or {}).get("repr") + level_name = (record.get("level") or {}).get("name") + message = record.get("message") + if not all((time_repr, level_name, message)): + continue + logs.append(f"{time_repr} - {level_name} - {message} - {extra}") + if len(logs) >= max_lines: + break + except (json.JSONDecodeError, AttributeError): + continue + + return logs[::-1] + + +__all__ = [ + "app_log_file", + "iter_file_lines_reversed", + "collect_task_logs", + "APP_LOG_FILENAME", + "DEFAULT_BLOCK_SIZE", + "MAX_TASK_LOG_LINES", +] diff --git a/openrag/core/utils/test_log_tail.py b/openrag/core/utils/test_log_tail.py new file mode 100644 index 000000000..d1eab3cd8 --- /dev/null +++ b/openrag/core/utils/test_log_tail.py @@ -0,0 +1,35 @@ +"""Tests for the reverse block file reader.""" + +from __future__ import annotations + +from pathlib import Path + +from core.utils.log_tail import APP_LOG_FILENAME, app_log_file, iter_file_lines_reversed + + +def test_app_log_file_under_configured_dir(): + assert app_log_file("/var/log/openrag") == Path("/var/log/openrag") / APP_LOG_FILENAME + + +def test_app_log_file_falls_back_to_logs_when_unset(): + assert app_log_file(None) == Path("logs") / APP_LOG_FILENAME + assert app_log_file("") == Path("logs") / APP_LOG_FILENAME + + +def test_reads_lines_newest_first_across_blocks(tmp_path): + f = tmp_path / "log.txt" + f.write_text("l1\nl2\nl3\nl4") + # small block size forces multi-block reads + assert list(iter_file_lines_reversed(f, block_size=4)) == ["l4", "l3", "l2", "l1"] + + +def test_skips_blank_lines(tmp_path): + f = tmp_path / "log.txt" + f.write_text("a\n\n\nb\n") + assert list(iter_file_lines_reversed(f)) == ["b", "a"] + + +def test_empty_file(tmp_path): + f = tmp_path / "log.txt" + f.write_text("") + assert list(iter_file_lines_reversed(f)) == [] diff --git a/openrag/core/utils/test_url_safety.py b/openrag/core/utils/test_url_safety.py new file mode 100644 index 000000000..6440bc70d --- /dev/null +++ b/openrag/core/utils/test_url_safety.py @@ -0,0 +1,37 @@ +"""Tests for the SSRF URL guard.""" + +from __future__ import annotations + +import pytest +from core.utils.url_safety import is_safe_url + + +@pytest.mark.parametrize( + "url", + [ + "http://127.0.0.1/x", + "https://localhost/x", + "http://10.0.0.1/x", + "http://192.168.1.1/x", + "http://169.254.169.254/latest/meta-data", # cloud metadata + "http://[::1]/x", # IPv6 loopback + "http://2130706433/x", # decimal-encoded 127.0.0.1 + "ftp://example.com/x", # non-http scheme + "file:///etc/passwd", + "not a url", + ], +) +def test_blocks_unsafe(url): + assert is_safe_url(url) is False + + +@pytest.mark.parametrize( + "url", + [ + "https://example.com/doc.pdf", + "http://files.internal.example.org/a/b.txt", + "https://8.8.8.8/x", # public IP literal + ], +) +def test_allows_public(url): + assert is_safe_url(url) is True diff --git a/openrag/core/utils/url_safety.py b/openrag/core/utils/url_safety.py new file mode 100644 index 000000000..c50d0608f --- /dev/null +++ b/openrag/core/utils/url_safety.py @@ -0,0 +1,75 @@ +"""SSRF guard for server-side URL fetches. + +Pure-stdlib host/address checks shared by any server-side fetcher (web-search +content fetch, MCP ``index_url``). Blocks loopback / private / link-local / +reserved / non-global addresses and the decimal-integer IPv4 encoding that +resolvers accept. Regular hostnames pass the literal check; callers that follow +redirects MUST re-validate every hop, since a public hostname can redirect to a +private target. +""" + +from __future__ import annotations + +import ipaddress +from urllib.parse import urlparse + + +def is_blocked_address(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + """True for any IP a server-side fetcher must not contact. + + Checks every private/reserved/non-global flag explicitly so the guard is + correct across Python minor releases (``is_global`` semantics shifted + between 3.10 and 3.11 for CGNAT and some multicast ranges). + """ + return ( + addr.is_loopback + or addr.is_private + or addr.is_link_local + or addr.is_reserved + or addr.is_unspecified + or addr.is_multicast + or not addr.is_global + ) + + +def is_safe_url(url: str) -> bool: + """Return True only if *url* is safe for a server-side fetch. + + Blocks non-HTTP(S) schemes, ``localhost``, IPv4/IPv6 literals in + private/loopback/link-local/reserved ranges, and decimal-integer-encoded + IPv4 (e.g. ``2130706433`` == ``127.0.0.1``). Regular hostnames pass — the + caller must re-check each redirect hop. + """ + try: + parsed = urlparse(url) + except Exception: + return False + + if parsed.scheme not in ("http", "https"): + return False + + host = parsed.hostname + if not host: + return False + + if host.lower() == "localhost": + return False + + # Dotted-decimal or IPv6 literal (e.g. "127.0.0.1", "::1", "10.0.0.1"). + try: + return not is_blocked_address(ipaddress.ip_address(host)) + except ValueError: + pass + + # Decimal-integer form (e.g. 2130706433 → 127.0.0.1). ip_address(int) + # interprets the value as a packed IPv4 address, matching glibc's resolver. + try: + return not is_blocked_address(ipaddress.ip_address(int(host))) + except (ValueError, TypeError): + pass + + # Regular hostname — passes the literal check; redirect hops re-validated. + return True + + +__all__ = ["is_blocked_address", "is_safe_url"] diff --git a/openrag/di/container.py b/openrag/di/container.py index d48622b50..3aa3022f8 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -59,6 +59,7 @@ from services.orchestrators.conversion_service import ConversionService from services.orchestrators.indexing_service import IndexingService from services.orchestrators.job_service import JobService + from services.orchestrators.mcp_service import MCPService from services.orchestrators.partition_service import PartitionService from services.orchestrators.query_service import QueryService from services.orchestrators.retrieval_service import RetrievalService @@ -106,6 +107,7 @@ def __init__(self, settings: Settings | None = None) -> None: self._indexing_service: IndexingService | None = None self._job_service: JobService | None = None self._conversion_service: ConversionService | None = None + self._mcp_service: MCPService | None = None def _require_settings(self) -> Settings: """Settings guard for the settings-dependent service properties. @@ -473,6 +475,36 @@ def conversion_service(self) -> ConversionService: ) return self._conversion_service + @property + def mcp_service(self) -> MCPService: + """MCPService — lazily built, cached for the container's lifetime. + + Composes the retrieval/partition/indexing/job/conversion + orchestrators plus the vector-store port into the application layer + the standalone MCP server (``api/mcp``) drives. Search defaults and + bounds come from the ``mcp`` settings section. + """ + if self._mcp_service is None: + from services.orchestrators.mcp_service import MCPService + + settings = self._require_settings() + mcp_cfg = settings.mcp + self._mcp_service = MCPService( + retrieval_service=self.retrieval_service, + partition_service=self.partition_service, + indexing_service=self.indexing_service, + job_service=self.job_service, + conversion_service=self.conversion_service, + vector_store=self.vector_store, + collection=settings.vectordb.collection_name, + default_top_k=mcp_cfg.default_top_k, + max_top_k=mcp_cfg.max_top_k, + similarity_threshold=mcp_cfg.similarity_threshold, + download_timeout=mcp_cfg.download_timeout, + max_download_bytes=mcp_cfg.max_download_bytes, + ) + return self._mcp_service + # ------------------------------------------------------------------ # Registry-based inference factories (Phase 6) # ------------------------------------------------------------------ diff --git a/openrag/di/providers.py b/openrag/di/providers.py index 1e6300a51..67c19ba5e 100644 --- a/openrag/di/providers.py +++ b/openrag/di/providers.py @@ -21,6 +21,7 @@ from services.orchestrators.conversion_service import ConversionService from services.orchestrators.indexing_service import IndexingService from services.orchestrators.job_service import JobService + from services.orchestrators.mcp_service import MCPService from services.orchestrators.partition_service import PartitionService from services.orchestrators.query_service import QueryService from services.orchestrators.retrieval_service import RetrievalService @@ -122,6 +123,11 @@ def get_conversion_service(request: Request = None) -> ConversionService: return _require_initialized(request).conversion_service +def get_mcp_service(request: Request = None) -> MCPService: + """Resolve the MCP orchestrator from the active container.""" + return _require_initialized(request).mcp_service + + def get_config(request: Request = None): """Resolve application configuration from the active container.""" return _require_initialized(request).config @@ -134,6 +140,7 @@ def get_config(request: Request = None): "get_conversion_service", "get_indexing_service", "get_job_service", + "get_mcp_service", "get_partition_service", "get_query_service", "get_retrieval_service", diff --git a/openrag/di/test_container.py b/openrag/di/test_container.py index 95678f5f0..bae9da2c6 100644 --- a/openrag/di/test_container.py +++ b/openrag/di/test_container.py @@ -237,11 +237,12 @@ def test_does_not_mutate_input_settings(self): ("indexing_service", "get_indexing_service"), ("job_service", "get_job_service"), ("conversion_service", "get_conversion_service"), + ("mcp_service", "get_mcp_service"), ] class TestPhase8OrchestratorWiring: - """8F: all nine orchestrators wired consistently (container + providers).""" + """8F: all orchestrators wired consistently (container + providers).""" @pytest.mark.parametrize("prop,_provider", _ORCHESTRATORS) def test_property_is_lazy_and_cache_slot_starts_none(self, prop, _provider): diff --git a/openrag/services/orchestrators/mcp_service.py b/openrag/services/orchestrators/mcp_service.py new file mode 100644 index 000000000..ef6a21b31 --- /dev/null +++ b/openrag/services/orchestrators/mcp_service.py @@ -0,0 +1,673 @@ +"""MCPService — application logic for the Model Context Protocol server. + +Reimplementation of PR 273's ``components/mcp`` (``SearchToolService`` + +``IndexationService``) on the hexagonal stack. Where the old code talked +straight to the Ray ``vectordb`` / ``indexer`` / ``task_state_manager`` +actors, this service composes the Phase-8 orchestrators +(``RetrievalService``, ``PartitionService``, ``IndexingService``, +``JobService``, ``ConversionService``) plus the :class:`VectorStore` port, +so it stays Ray-free (``services → core`` only). + +Per-request authorization (the OpenRAG partition ACL) is passed in +explicitly — ``user_id`` / ``is_admin`` / ``allowed_partitions`` are method +arguments, never read from ambient context. The ``api/mcp`` server layer +resolves them from the bearer token and forwards them; this keeps the +service a pure orchestrator that is trivially unit-testable. + +``allowed_partitions == ["all"]`` is the admin wildcard (set by the server +when ``is_admin and SUPER_ADMIN_MODE``); ``None`` means the auth context +never ran and every method refuses. +""" + +from __future__ import annotations + +import asyncio +import difflib +import ipaddress +import json +import mimetypes +import socket +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING, Any +from urllib.parse import urljoin, urlparse + +import httpx +from core.utils.log_tail import collect_task_logs +from core.utils.url_safety import is_blocked_address, is_safe_url +from utils.logger import get_logger + +if TYPE_CHECKING: + from core.vector_stores import VectorStore + from services.orchestrators.conversion_service import ConversionService + from services.orchestrators.indexing_service import IndexingService + from services.orchestrators.job_service import JobService + from services.orchestrators.partition_service import PartitionService + from services.orchestrators.retrieval_service import RetrievalService + +logger = get_logger() + +_ROLE_HIERARCHY: dict[str, int] = {"viewer": 1, "editor": 2, "owner": 3} +_FORBIDDEN_FILE_ID_CHARS: frozenset[str] = frozenset("/") +_ACTIVE_STATES = {"QUEUED", "SERIALIZING", "CHUNKING", "INSERTING"} +_MAX_REDIRECTS = 10 +_MAX_CHUNKS_PER_CALL = 200 +_MAX_FUZZY_CANDIDATES = 5000 + +# Catalog/store-managed fields a caller must not be able to spoof via the +# free-form metadata dict on write tools. ``partition`` is intentionally NOT +# here — update_file_metadata uses it as an authorized move control (and +# re-checks editor access on the destination). +_PROTECTED_METADATA_KEYS: frozenset[str] = frozenset( + {"file_id", "source", "created_by", "file_size", "file_count", "_id", "vector", "text"} +) + + +def _strip_protected_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]: + """Drop server-managed keys from caller-supplied metadata (defense in depth).""" + md = dict(metadata or {}) + removed = [k for k in md if k in _PROTECTED_METADATA_KEYS] + for k in removed: + del md[k] + if removed: + logger.warning("Dropped protected metadata keys from MCP write", keys=sorted(removed)) + return md + + +class MCPService: + """Search + catalog + indexation operations exposed as MCP tools.""" + + def __init__( + self, + *, + retrieval_service: RetrievalService, + partition_service: PartitionService, + indexing_service: IndexingService, + job_service: JobService, + conversion_service: ConversionService, + vector_store: VectorStore, + collection: str, + default_top_k: int = 5, + max_top_k: int = 50, + similarity_threshold: float = 0.8, + download_timeout: float = 30.0, + max_download_bytes: int = 100 * 1024 * 1024, + ) -> None: + self._retrieval = retrieval_service + self._partitions = partition_service + self._indexing = indexing_service + self._jobs = job_service + self._conversion = conversion_service + self._vector_store = vector_store + self._collection = collection + self.default_top_k = default_top_k + self.max_top_k = max_top_k + self.similarity_threshold = similarity_threshold + self.download_timeout = download_timeout + self.max_download_bytes = max_download_bytes + + # ------------------------------------------------------------------ + # Authorization helpers + # ------------------------------------------------------------------ + + @staticmethod + def _validate_file_id(file_id: str) -> None: + if not file_id or not file_id.strip(): + raise ValueError("file_id cannot be empty") + bad = _FORBIDDEN_FILE_ID_CHARS & set(file_id) + if bad: + raise ValueError(f"file_id contains forbidden characters: {', '.join(sorted(bad))}") + + @staticmethod + def _enforce_partition_scope( + requested_partitions: list[str], + allowed_partitions: list[str] | None, + ) -> list[str]: + """Resolve the search scope against the caller's allowed partitions.""" + if allowed_partitions is None: + raise PermissionError("Authentication context is missing") + if allowed_partitions == ["all"]: + return ["all"] if requested_partitions == ["all"] else requested_partitions + if requested_partitions == ["all"]: + return allowed_partitions + denied = [p for p in requested_partitions if p not in allowed_partitions] + if denied: + raise PermissionError(f"Access denied for partition(s): {', '.join(sorted(set(denied)))}") + return requested_partitions + + @staticmethod + def _enforce_partition_access(partition: str, allowed_partitions: list[str] | None) -> None: + if allowed_partitions is None: + raise PermissionError("Authentication context is missing") + if allowed_partitions == ["all"]: + return + if partition not in allowed_partitions: + raise PermissionError(f"Access denied for partition: {partition}") + + async def _enforce_editor_access( + self, + partition: str, + allowed_partitions: list[str] | None, + user_id: int | None, + ) -> None: + """Require editor (or owner) role on *partition*.""" + self._enforce_partition_access(partition, allowed_partitions) + if allowed_partitions == ["all"]: + return # admin wildcard + if user_id is None: + return # no-auth mode — treated as admin + members = await self._partitions.list_members(partition) + membership = next((m for m in members if m.get("user_id") == user_id), None) + if membership is None or _ROLE_HIERARCHY.get(membership.get("role", ""), 0) < _ROLE_HIERARCHY["editor"]: + raise PermissionError(f"Editor role required for partition: {partition}") + + async def _ensure_partition_exists( + self, + partition: str, + allowed_partitions: list[str] | None, + user_id: int | None, + ) -> list[str] | None: + """Auto-create *partition* (owned by the caller) when it is missing. + + Returns the possibly-extended allowed list. The original list is + never mutated — it may be shared across concurrent requests. + """ + if allowed_partitions is None: + return None # _enforce_*_access will raise + if allowed_partitions == ["all"]: + return allowed_partitions + if partition in allowed_partitions: + return allowed_partitions + if await self._partitions.partition_exists(partition): + return allowed_partitions # exists but no membership → access check raises + await self._partitions.create_partition(partition, user_id) + return [*allowed_partitions, partition] + + # ------------------------------------------------------------------ + # Search + # ------------------------------------------------------------------ + + async def search_documents( + self, + *, + query: str, + partitions: list[str] | None, + top_k: int | None, + allowed_partitions: list[str] | None, + file_id: str | None = None, + ) -> dict[str, Any]: + normalized_query = query.strip() + if not normalized_query: + raise ValueError("Query cannot be empty") + if file_id is not None: + self._validate_file_id(file_id) + + requested = [p.strip() for p in (partitions or ["all"]) if p and p.strip()] or ["all"] + scoped = self._enforce_partition_scope(requested, allowed_partitions) + + effective_top_k = top_k if top_k is not None else self.default_top_k + if effective_top_k <= 0: + raise ValueError("top_k must be greater than 0") + effective_top_k = min(effective_top_k, self.max_top_k) + + # Inline the file_id as a literal Milvus expression rather than using + # the templated ``filter_params`` form: the shared VectorStoreSearcher + # forwards the raw ``filter`` string to Milvus but drops + # ``filter_params``, so a ``{var}`` placeholder would never be bound. + # ``json.dumps`` yields a correctly quoted/escaped string literal. + search_filter = f"file_id == {json.dumps(file_id)}" if file_id is not None else None + + chunks = await self._retrieval.search( + text=normalized_query, + partitions=scoped, + top_k=effective_top_k, + similarity_threshold=self.similarity_threshold, + filter=search_filter, + ) + documents = [self._shape_chunk(c) for c in chunks] + return { + "query": normalized_query, + "partitions": scoped, + "top_k": effective_top_k, + "count": len(documents), + "documents": documents, + } + + @staticmethod + def _shape_chunk(chunk: Any) -> dict[str, Any]: + """Domain ``Chunk`` → MCP search-result dict (legacy metadata shape).""" + meta = chunk.to_langchain().metadata + return { + "chunk_id": meta.get("_id") or chunk.id, + "content": chunk.text, + "metadata": meta, + } + + # ------------------------------------------------------------------ + # Partitions & files + # ------------------------------------------------------------------ + + async def list_partitions(self, *, allowed_partitions: list[str] | None) -> dict[str, Any]: + if allowed_partitions is None: + raise PermissionError("Authentication context is missing") + partitions = await self._partitions.list_partitions() + if allowed_partitions != ["all"]: + partitions = [p for p in partitions if p["partition"] in allowed_partitions] + return {"count": len(partitions), "partitions": partitions} + + async def list_files( + self, + *, + partition: str, + allowed_partitions: list[str] | None, + limit: int | None = None, + ) -> dict[str, Any]: + self._enforce_partition_access(partition, allowed_partitions) + files = await self._partitions.list_files(partition, limit=limit) + return {"partition": partition, "count": len(files), "files": files} + + async def get_file_info( + self, + *, + partition: str, + file_id: str, + allowed_partitions: list[str] | None, + ) -> dict[str, Any]: + self._enforce_partition_access(partition, allowed_partitions) + if not await self._partitions.file_exists(file_id, partition): + raise FileNotFoundError(f"File '{file_id}' not found in partition '{partition}'") + # Read straight from the vector store (same source as get_file_chunks) so + # chunk_count is exact — PartitionService.get_file_chunks caps at 2000, + # which would silently undercount large files and disagree with the + # total_chunks reported by get_file_chunks. + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + output_fields=["*"], + ) + metadata = {k: v for k, v in rows[0].items() if k not in ("_id", "text", "vector")} if rows else {} + return { + "partition": partition, + "file_id": file_id, + "chunk_count": len(rows), + "metadata": metadata, + } + + async def get_file_chunks( + self, + *, + partition: str, + file_id: str, + allowed_partitions: list[str] | None, + offset: int = 0, + limit: int = 3, + ) -> dict[str, Any]: + self._enforce_partition_access(partition, allowed_partitions) + if offset < 0: + raise ValueError("offset must be >= 0") + if limit == 0 or limit < -1: + raise ValueError("limit must be -1 (all) or a positive integer") + # Cap the page size so a single call can never return an unbounded + # number of (potentially large) chunks. -1 means "as many as allowed". + page_size = _MAX_CHUNKS_PER_CALL if limit == -1 else min(limit, _MAX_CHUNKS_PER_CALL) + if not await self._partitions.file_exists(file_id, partition): + raise FileNotFoundError(f"File '{file_id}' not found in partition '{partition}'") + rows = await self._vector_store.query_chunks_by_filter( + self._collection, + {"partition": partition, "file_id": file_id}, + output_fields=["*"], + ) + # Milvus _id is INT64; guard against missing/non-int values so a + # malformed row can't raise a TypeError mid-sort. + rows.sort(key=lambda r: r["_id"] if isinstance(r.get("_id"), int) else 0) + total = len(rows) + page = rows[offset : offset + page_size] + return { + "partition": partition, + "file_id": file_id, + "total_chunks": total, + "offset": offset, + "limit": limit, + "has_more": offset + len(page) < total, + "chunks": [ + { + "chunk_id": row.get("_id"), + "content": row.get("text"), + "metadata": {k: v for k, v in row.items() if k not in ("text", "_id", "vector")}, + } + for row in page + ], + } + + async def fuzzy_search_files( + self, + *, + query: str, + allowed_partitions: list[str] | None, + partition: str | None = None, + cutoff: float = 0.4, + limit: int = 20, + ) -> dict[str, Any]: + if allowed_partitions is None: + raise PermissionError("Authentication context is missing") + normalized_query = query.strip().lower() + if not normalized_query: + raise ValueError("Query cannot be empty") + + if partition is not None: + self._enforce_partition_access(partition, allowed_partitions) + search_partitions = [partition] + elif allowed_partitions == ["all"]: + search_partitions = [p["partition"] for p in await self._partitions.list_partitions()] + else: + search_partitions = allowed_partitions + + # Bound the candidate set so an admin (or a user with many partitions) + # can't trigger an unbounded scan + difflib pass over the whole catalog. + candidate_files: list[dict[str, Any]] = [] + for part in search_partitions: + if len(candidate_files) >= _MAX_FUZZY_CANDIDATES: + logger.warning("fuzzy_search_files candidate cap reached", cap=_MAX_FUZZY_CANDIDATES) + break + candidate_files.extend(await self._partitions.list_files(part)) + candidate_files = candidate_files[:_MAX_FUZZY_CANDIDATES] + + scored: list[tuple[float, dict[str, Any]]] = [] + for f in candidate_files: + names = [str(f[k]).lower() for k in ("filename", "original_filename", "file_id") if f.get(k)] + if not names: + continue + best = max(difflib.SequenceMatcher(None, normalized_query, name).ratio() for name in names) + if best >= cutoff: + scored.append((best, f)) + + scored.sort(key=lambda x: x[0], reverse=True) + results = [{"score": round(ratio, 4), **f} for ratio, f in scored[:limit]] + return {"query": query, "count": len(results), "files": results} + + # ------------------------------------------------------------------ + # Tasks + # ------------------------------------------------------------------ + + async def get_task_status( + self, + *, + task_id: str, + user_id: int | None, + is_admin: bool, + ) -> dict[str, Any]: + details = await self._jobs.get_task_details(task_id) + if details is None: + raise KeyError(f"Task '{task_id}' not found") + if not is_admin and user_id is not None and details.get("user_id") != user_id: + raise PermissionError("You do not have permission to access this task") + + state = await self._indexing.get_task_state(task_id) + result: dict[str, Any] = {"task_id": task_id, "task_state": state, "details": details} + if state == "FAILED": + result["error"] = await self._indexing.get_task_error(task_id) + return result + + async def list_my_tasks( + self, + *, + user_id: int | None, + is_admin: bool, + task_status: str | None = None, + ) -> dict[str, Any]: + tasks = await self._jobs.list_tasks(is_admin=is_admin, user_id=user_id, task_status=task_status) + for task in tasks: + if task.get("state") == "FAILED": + task["error"] = await self._indexing.get_task_error(task["task_id"]) + return {"count": len(tasks), "tasks": tasks} + + async def get_task_logs( + self, + *, + task_id: str, + user_id: int | None, + is_admin: bool, + log_file: str | Path, + max_lines: int = 100, + ) -> dict[str, Any]: + details = await self._jobs.get_task_details(task_id) + if details is None: + raise KeyError(f"Task '{task_id}' not found") + if not is_admin and user_id is not None and details.get("user_id") != user_id: + raise PermissionError("You do not have permission to access this task") + + log_path = Path(log_file) + if not log_path.exists(): + raise FileNotFoundError(f"Log file not found: {log_path}") + logs = collect_task_logs(log_path, task_id, max_lines) + return {"task_id": task_id, "count": len(logs), "logs": logs} + + # ------------------------------------------------------------------ + # Chunk lookup + # ------------------------------------------------------------------ + + async def get_chunk_by_id( + self, + *, + chunk_id: str, + allowed_partitions: list[str] | None, + ) -> dict[str, Any]: + if allowed_partitions is None: + raise PermissionError("Authentication context is missing") + chunk = await self._conversion.get_chunk(chunk_id) + if chunk is None: + raise KeyError(f"Chunk '{chunk_id}' not found") + chunk_partition = chunk["metadata"].get("partition") + if allowed_partitions != ["all"] and chunk_partition not in allowed_partitions: + raise PermissionError(f"Access denied for partition: {chunk_partition}") + return { + "chunk_id": chunk_id, + "page_content": chunk["page_content"], + "metadata": chunk["metadata"], + } + + # ------------------------------------------------------------------ + # Write operations + # ------------------------------------------------------------------ + + async def delete_file( + self, + *, + partition: str, + file_id: str, + allowed_partitions: list[str] | None, + user_id: int | None, + ) -> dict[str, Any]: + self._validate_file_id(file_id) + await self._enforce_editor_access(partition, allowed_partitions, user_id) + if not await self._partitions.file_exists(file_id, partition): + raise FileNotFoundError(f"File '{file_id}' not found in partition '{partition}'") + await self._indexing.delete_file(file_id, partition) + return { + "partition": partition, + "file_id": file_id, + "message": f"File '{file_id}' deleted from partition '{partition}'.", + } + + async def update_file_metadata( + self, + *, + partition: str, + file_id: str, + metadata: dict[str, Any], + allowed_partitions: list[str] | None, + user_id: int | None, + ) -> dict[str, Any]: + self._validate_file_id(file_id) + await self._enforce_editor_access(partition, allowed_partitions, user_id) + # ``partition`` in the payload is an authorized move control (not a + # protected key) — require editor on the destination too. + if metadata and "partition" in metadata: + await self._enforce_editor_access(metadata["partition"], allowed_partitions, user_id) + if not await self._partitions.file_exists(file_id, partition): + raise FileNotFoundError(f"File '{file_id}' not found in partition '{partition}'") + await self._indexing.update_metadata( + file_id, + _strip_protected_metadata(metadata), + partition, + {"id": user_id} if user_id is not None else None, + ) + return { + "partition": partition, + "file_id": file_id, + "message": f"Metadata for file '{file_id}' successfully updated.", + } + + async def copy_file( + self, + *, + source_partition: str, + source_file_id: str, + dest_partition: str, + dest_file_id: str, + allowed_partitions: list[str] | None, + user_id: int | None, + extra_metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + self._validate_file_id(source_file_id) + self._validate_file_id(dest_file_id) + self._enforce_partition_access(source_partition, allowed_partitions) + await self._enforce_editor_access(dest_partition, allowed_partitions, user_id) + + if not await self._partitions.file_exists(source_file_id, source_partition): + raise FileNotFoundError(f"File '{source_file_id}' not found in partition '{source_partition}'") + if await self._partitions.file_exists(dest_file_id, dest_partition): + raise FileExistsError(f"File '{dest_file_id}' already exists in partition '{dest_partition}'") + + await self._indexing.copy_file( + source_file_id=source_file_id, + source_partition=source_partition, + target_file_id=dest_file_id, + target_partition=dest_partition, + metadata=_strip_protected_metadata(extra_metadata), + user={"id": user_id} if user_id is not None else None, + ) + return { + "source_partition": source_partition, + "source_file_id": source_file_id, + "dest_partition": dest_partition, + "dest_file_id": dest_file_id, + "message": "File copied successfully.", + } + + async def index_url( + self, + *, + url: str, + partition: str, + file_id: str, + allowed_partitions: list[str] | None, + user_id: int | None, + extra_metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + # Validate cheap/safe inputs BEFORE any side effect (partition + # auto-create), so a blocked URL or bad id never leaves an empty + # partition behind. + self._validate_file_id(file_id) + # SSRF guard: reject non-http(s) and any host resolving to a + # loopback/private/link-local/reserved range (cloud metadata, internal + # services). Redirect hops are re-validated inside _safe_download. + if not is_safe_url(url): + raise ValueError(f"URL is not allowed for server-side fetch: {url!r}") + + allowed_partitions = await self._ensure_partition_exists(partition, allowed_partitions, user_id) + await self._enforce_editor_access(partition, allowed_partitions, user_id) + + if await self._partitions.file_exists(file_id, partition): + raise FileExistsError(f"File '{file_id}' already exists in partition '{partition}'") + + filename = Path(urlparse(url).path.rstrip("/")).name or file_id + suffix = Path(filename).suffix or "" + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp_path = Path(tmp.name) + try: + await self._safe_download(url, tmp_path) + except Exception as exc: + tmp_path.unlink(missing_ok=True) + raise RuntimeError(f"Failed to download '{url}': {exc}") from exc + + metadata = _strip_protected_metadata(extra_metadata) + metadata["source_url"] = url + guessed_mime, _ = mimetypes.guess_type(filename) + if guessed_mime and "mimetype" not in metadata: + metadata["mimetype"] = guessed_mime + + task_id = await self._indexing.add_file( + file_path=str(tmp_path), + file_id=file_id, + partition=partition, + metadata=metadata, + sanitized_filename=filename, + original_filename=filename, + user={"id": user_id} if user_id is not None else None, + ) + return { + "partition": partition, + "file_id": file_id, + "task_id": task_id, + "message": f"Indexation started. Poll get_indexation_task_status with task_id='{task_id}'.", + } + + @staticmethod + async def _assert_host_resolves_safely(host: str) -> None: + """Reject *host* if it resolves to any private/loopback/reserved address. + + ``is_safe_url`` only inspects the literal host, so a public hostname + pointing at an internal IP (DNS-based SSRF, incl. cloud metadata at + 169.254.169.254) would otherwise pass. Resolution runs on the event + loop's async resolver so the loop is never blocked. This narrows but + cannot fully close TOCTOU DNS rebinding — the literal check and the + per-hop re-validation in ``_safe_download`` remain the other layers. + """ + if not host: + raise ValueError("URL has no host") + loop = asyncio.get_running_loop() + try: + infos = await loop.getaddrinfo(host, None, type=socket.SOCK_STREAM) + except socket.gaierror as exc: + raise ValueError(f"could not resolve host {host!r}: {exc}") from exc + for info in infos: + addr = ipaddress.ip_address(info[4][0]) + if is_blocked_address(addr): + raise ValueError(f"host {host!r} resolves to a disallowed address: {addr}") + + async def _safe_download(self, url: str, dest: Path) -> None: + """Stream *url* to *dest* with SSRF, redirect, size and time guards. + + Redirects are followed manually (``follow_redirects=False``) so every + hop is re-validated by :func:`is_safe_url` — a public host can otherwise + redirect to a private address after the initial check. The body is + streamed with a hard byte ceiling so a hostile endpoint cannot exhaust + disk, and the whole fetch is bounded by ``download_timeout``. + """ + current_url = url + async with httpx.AsyncClient(timeout=self.download_timeout, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + if not is_safe_url(current_url): + raise ValueError(f"redirect to a disallowed URL: {current_url!r}") + await self._assert_host_resolves_safely(urlparse(current_url).hostname or "") + async with client.stream("GET", current_url) as response: + if response.is_redirect: + location = response.headers.get("location", "") + if not location: + raise ValueError("redirect without a Location header") + current_url = urljoin(current_url, location) + continue + response.raise_for_status() + written = 0 + with dest.open("wb") as fh: + async for chunk in response.aiter_bytes(): + written += len(chunk) + if written > self.max_download_bytes: + raise ValueError(f"download exceeds {self.max_download_bytes} bytes") + fh.write(chunk) + return + raise ValueError("too many redirects") + + +__all__ = ["MCPService"] diff --git a/openrag/services/orchestrators/test_mcp_service.py b/openrag/services/orchestrators/test_mcp_service.py new file mode 100644 index 000000000..609702e81 --- /dev/null +++ b/openrag/services/orchestrators/test_mcp_service.py @@ -0,0 +1,738 @@ +"""Unit tests for :class:`MCPService` — the MCP application orchestrator. + +The orchestrators it composes are replaced by small fakes so these tests +exercise the ACL scope/role enforcement, response shaping, pagination, +fuzzy ranking, task assembly and URL-indexation guards in isolation. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import httpx +import pytest +from services.orchestrators.mcp_service import MCPService + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class FakeChunk: + def __init__(self, *, _id, text, metadata=None): + self.id = str(_id) + self.text = text + self._meta = {"_id": _id, **(metadata or {})} + + def to_langchain(self, *, with_id: bool = True): + return SimpleNamespace(metadata=dict(self._meta)) + + +class FakeRetrieval: + def __init__(self, *, chunks=None): + self._chunks = chunks if chunks is not None else [] + self.calls: list[dict] = [] + + async def search(self, **kwargs): + self.calls.append(kwargs) + return list(self._chunks) + + +class FakePartitions: + def __init__(self, *, partitions=None, files=None, exists=True, chunks=None, members=None, partition_exists=True): + self._partitions = partitions if partitions is not None else [] + self._files = files if files is not None else {} + self._exists = exists + self._chunks = chunks if chunks is not None else [] + self._members = members if members is not None else [] + self._partition_exists = partition_exists + self.created: list[tuple[str, int | None]] = [] + + async def list_partitions(self): + return list(self._partitions) + + async def list_files(self, partition, limit=None): + files = self._files.get(partition, []) if isinstance(self._files, dict) else list(self._files) + return files[:limit] if limit is not None else files + + async def file_exists(self, file_id, partition): + return self._exists + + async def get_file_chunks(self, partition, file_id, limit=2000): + return list(self._chunks) + + async def list_members(self, partition): + return list(self._members) + + async def partition_exists(self, partition): + return self._partition_exists + + async def create_partition(self, partition, user_id): + # Mirror PgPartitionRepository: the creator is granted owner. + self.created.append((partition, user_id)) + self._members.append({"user_id": user_id, "role": "owner"}) + + +class FakeIndexing: + def __init__(self, *, state="COMPLETED", error="boom"): + self._state = state + self._error = error + self.deleted: list[tuple[str, str]] = [] + self.updated: list[tuple] = [] + self.copied: list[dict] = [] + self.added: list[dict] = [] + + async def get_task_state(self, task_id): + return self._state + + async def get_task_error(self, task_id): + return self._error + + async def delete_file(self, file_id, partition): + self.deleted.append((file_id, partition)) + + async def update_metadata(self, file_id, metadata, partition, user): + self.updated.append((file_id, metadata, partition, user)) + + async def copy_file(self, **kwargs): + self.copied.append(kwargs) + + async def add_file(self, **kwargs): + self.added.append(kwargs) + return "task-123" + + +class FakeJobs: + def __init__(self, *, details=None, tasks=None): + self._details = details + self._tasks = tasks if tasks is not None else [] + + async def get_task_details(self, task_id): + return self._details + + async def list_tasks(self, *, is_admin, user_id, task_status=None): + self.last = {"is_admin": is_admin, "user_id": user_id, "task_status": task_status} + return [dict(t) for t in self._tasks] + + +class FakeConversion: + def __init__(self, *, chunk=None): + self._chunk = chunk + + async def get_chunk(self, chunk_id): + return self._chunk + + +class FakeVectorStore: + def __init__(self, *, rows=None): + self._rows = rows if rows is not None else [] + self.queries: list[tuple] = [] + + async def query_chunks_by_filter(self, collection, filters, output_fields=None): + self.queries.append((collection, filters, output_fields)) + return [dict(r) for r in self._rows] + + +def _service( + *, + retrieval=None, + partitions=None, + indexing=None, + jobs=None, + conversion=None, + vector_store=None, + default_top_k=5, + max_top_k=50, + similarity_threshold=0.8, +): + return MCPService( + retrieval_service=retrieval or FakeRetrieval(), + partition_service=partitions or FakePartitions(), + indexing_service=indexing or FakeIndexing(), + job_service=jobs or FakeJobs(), + conversion_service=conversion or FakeConversion(), + vector_store=vector_store or FakeVectorStore(), + collection="chunks", + default_top_k=default_top_k, + max_top_k=max_top_k, + similarity_threshold=similarity_threshold, + ) + + +# --------------------------------------------------------------------------- +# Search + scope ACL +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_search_empty_query_raises(): + with pytest.raises(ValueError): + await _service().search_documents(query=" ", partitions=None, top_k=None, allowed_partitions=["all"]) + + +@pytest.mark.asyncio +async def test_search_missing_auth_context_raises(): + with pytest.raises(PermissionError): + await _service().search_documents(query="hi", partitions=None, top_k=None, allowed_partitions=None) + + +@pytest.mark.asyncio +async def test_search_admin_wildcard_passthrough(): + retrieval = FakeRetrieval() + await _service(retrieval=retrieval).search_documents( + query="hi", partitions=None, top_k=None, allowed_partitions=["all"] + ) + assert retrieval.calls[0]["partitions"] == ["all"] + + +@pytest.mark.asyncio +async def test_search_all_request_resolves_to_allowed(): + retrieval = FakeRetrieval() + await _service(retrieval=retrieval).search_documents( + query="hi", partitions=["all"], top_k=None, allowed_partitions=["a", "b"] + ) + assert retrieval.calls[0]["partitions"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_search_denied_partition_raises(): + with pytest.raises(PermissionError): + await _service().search_documents(query="hi", partitions=["c"], top_k=None, allowed_partitions=["a", "b"]) + + +@pytest.mark.asyncio +async def test_search_top_k_capped_to_max(): + retrieval = FakeRetrieval() + out = await _service(retrieval=retrieval, max_top_k=50).search_documents( + query="hi", partitions=["a"], top_k=999, allowed_partitions=["a"] + ) + assert retrieval.calls[0]["top_k"] == 50 + assert out["top_k"] == 50 + + +@pytest.mark.asyncio +async def test_search_top_k_default_used(): + retrieval = FakeRetrieval() + await _service(retrieval=retrieval, default_top_k=7).search_documents( + query="hi", partitions=["a"], top_k=None, allowed_partitions=["a"] + ) + assert retrieval.calls[0]["top_k"] == 7 + + +@pytest.mark.asyncio +async def test_search_nonpositive_top_k_raises(): + with pytest.raises(ValueError): + await _service().search_documents(query="hi", partitions=["a"], top_k=0, allowed_partitions=["a"]) + + +@pytest.mark.asyncio +async def test_search_file_id_builds_filter(): + retrieval = FakeRetrieval() + svc = _service(retrieval=retrieval) + await svc.search_documents(query="hi", partitions=["a"], top_k=3, allowed_partitions=["a"], file_id="f1") + call = retrieval.calls[0] + # Inlined as a literal expr (the shared searcher drops filter_params). + assert call["filter"] == 'file_id == "f1"' + assert "filter_params" not in call + + +@pytest.mark.asyncio +async def test_search_rejects_file_id_with_slash(): + with pytest.raises(ValueError): + await _service().search_documents( + query="hi", partitions=["a"], top_k=3, allowed_partitions=["a"], file_id="a/b" + ) + + +@pytest.mark.asyncio +async def test_search_shapes_chunks(): + chunks = [FakeChunk(_id=11, text="body", metadata={"file_id": "f1"})] + out = await _service(retrieval=FakeRetrieval(chunks=chunks)).search_documents( + query="hi", partitions=["a"], top_k=3, allowed_partitions=["a"] + ) + assert out["count"] == 1 + doc = out["documents"][0] + assert doc["chunk_id"] == 11 + assert doc["content"] == "body" + assert doc["metadata"]["file_id"] == "f1" + + +# --------------------------------------------------------------------------- +# Partitions & files +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_list_partitions_filtered_by_allowed(): + parts = [{"partition": "a"}, {"partition": "b"}, {"partition": "c"}] + out = await _service(partitions=FakePartitions(partitions=parts)).list_partitions(allowed_partitions=["a", "c"]) + assert {p["partition"] for p in out["partitions"]} == {"a", "c"} + assert out["count"] == 2 + + +@pytest.mark.asyncio +async def test_list_partitions_admin_sees_all(): + parts = [{"partition": "a"}, {"partition": "b"}] + out = await _service(partitions=FakePartitions(partitions=parts)).list_partitions(allowed_partitions=["all"]) + assert out["count"] == 2 + + +@pytest.mark.asyncio +async def test_list_files_denied_raises(): + with pytest.raises(PermissionError): + await _service().list_files(partition="x", allowed_partitions=["a"]) + + +@pytest.mark.asyncio +async def test_get_file_info_not_found(): + svc = _service(partitions=FakePartitions(exists=False)) + with pytest.raises(FileNotFoundError): + await svc.get_file_info(partition="a", file_id="missing", allowed_partitions=["a"]) + + +@pytest.mark.asyncio +async def test_get_file_info_strips_id_text_vector_from_metadata(): + rows = [{"_id": 1, "file_id": "f1", "page": 2, "text": "body", "vector": [0.1, 0.2]}] + svc = _service(partitions=FakePartitions(exists=True), vector_store=FakeVectorStore(rows=rows)) + out = await svc.get_file_info(partition="a", file_id="f1", allowed_partitions=["a"]) + assert out["chunk_count"] == 1 + assert {"_id", "text", "vector"}.isdisjoint(out["metadata"]) + assert out["metadata"]["file_id"] == "f1" + + +@pytest.mark.asyncio +async def test_get_file_info_count_not_capped(): + # Exact count from the vector store, not PartitionService's 2000 cap. + rows = [{"_id": i, "file_id": "f1"} for i in range(2500)] + svc = _service(partitions=FakePartitions(exists=True), vector_store=FakeVectorStore(rows=rows)) + out = await svc.get_file_info(partition="a", file_id="f1", allowed_partitions=["a"]) + assert out["chunk_count"] == 2500 + + +@pytest.mark.asyncio +async def test_get_file_chunks_paginates_with_content(): + rows = [{"_id": i, "text": f"c{i}", "file_id": "f1", "partition": "a"} for i in range(5)] + svc = _service(vector_store=FakeVectorStore(rows=rows)) + out = await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], offset=1, limit=2) + assert out["total_chunks"] == 5 + assert out["has_more"] is True + assert [c["chunk_id"] for c in out["chunks"]] == [1, 2] + assert out["chunks"][0]["content"] == "c1" + assert "text" not in out["chunks"][0]["metadata"] + + +@pytest.mark.asyncio +async def test_get_file_chunks_limit_all(): + rows = [{"_id": i, "text": f"c{i}"} for i in range(3)] + svc = _service(vector_store=FakeVectorStore(rows=rows)) + out = await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], offset=0, limit=-1) + assert len(out["chunks"]) == 3 + assert out["has_more"] is False + + +@pytest.mark.asyncio +async def test_fuzzy_search_ranks_and_cuts_off(): + files = { + "a": [{"file_id": "annual_report", "filename": "annual_report.pdf"}, {"file_id": "zzz", "filename": "zzz.txt"}] + } + svc = _service(partitions=FakePartitions(files=files)) + out = await svc.fuzzy_search_files(query="annual report", allowed_partitions=["a"], partition="a", cutoff=0.4) + assert out["count"] == 1 + assert out["files"][0]["file_id"] == "annual_report" + assert "score" in out["files"][0] + + +@pytest.mark.asyncio +async def test_fuzzy_search_empty_query_raises(): + with pytest.raises(ValueError): + await _service().fuzzy_search_files(query=" ", allowed_partitions=["a"]) + + +# --------------------------------------------------------------------------- +# Tasks +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_task_status_not_found(): + with pytest.raises(KeyError): + await _service(jobs=FakeJobs(details=None)).get_task_status(task_id="t1", user_id=1, is_admin=False) + + +@pytest.mark.asyncio +async def test_task_status_owner_mismatch_raises(): + jobs = FakeJobs(details={"user_id": 2}) + with pytest.raises(PermissionError): + await _service(jobs=jobs).get_task_status(task_id="t1", user_id=1, is_admin=False) + + +@pytest.mark.asyncio +async def test_task_status_admin_bypasses_owner(): + jobs = FakeJobs(details={"user_id": 2}) + out = await _service(jobs=jobs, indexing=FakeIndexing(state="COMPLETED")).get_task_status( + task_id="t1", user_id=1, is_admin=True + ) + assert out["task_state"] == "COMPLETED" + assert "error" not in out + + +@pytest.mark.asyncio +async def test_task_status_failed_includes_error(): + jobs = FakeJobs(details={"user_id": 1}) + out = await _service(jobs=jobs, indexing=FakeIndexing(state="FAILED", error="kaboom")).get_task_status( + task_id="t1", user_id=1, is_admin=False + ) + assert out["task_state"] == "FAILED" + assert out["error"] == "kaboom" + + +@pytest.mark.asyncio +async def test_list_my_tasks_forwards_scope_and_decorates_failed(): + jobs = FakeJobs( + tasks=[ + {"task_id": "t1", "state": "FAILED", "details": {}}, + {"task_id": "t2", "state": "COMPLETED", "details": {}}, + ] + ) + out = await _service(jobs=jobs, indexing=FakeIndexing(error="why")).list_my_tasks( + user_id=5, is_admin=False, task_status="failed" + ) + assert jobs.last == {"is_admin": False, "user_id": 5, "task_status": "failed"} + failed = next(t for t in out["tasks"] if t["task_id"] == "t1") + assert failed["error"] == "why" + completed = next(t for t in out["tasks"] if t["task_id"] == "t2") + assert "error" not in completed + + +@pytest.mark.asyncio +async def test_get_task_logs_ownership_and_parsing(tmp_path): + log = tmp_path / "app.json" + lines = [ + {"record": {"time": {"repr": "T1"}, "level": {"name": "INFO"}, "message": "first", "extra": {"task_id": "t1"}}}, + { + "record": { + "time": {"repr": "T2"}, + "level": {"name": "INFO"}, + "message": "other", + "extra": {"task_id": "zzz"}, + } + }, + { + "record": { + "time": {"repr": "T3"}, + "level": {"name": "ERROR"}, + "message": "second", + "extra": {"task_id": "t1"}, + } + }, + ] + log.write_text("\n".join(json.dumps(line) for line in lines)) + svc = _service(jobs=FakeJobs(details={"user_id": 1})) + out = await svc.get_task_logs(task_id="t1", user_id=1, is_admin=False, log_file=log, max_lines=100) + assert out["count"] == 2 + assert "first" in out["logs"][0] + assert "second" in out["logs"][1] + + +@pytest.mark.asyncio +async def test_get_task_logs_missing_file_raises(tmp_path): + svc = _service(jobs=FakeJobs(details={"user_id": 1})) + with pytest.raises(FileNotFoundError): + await svc.get_task_logs(task_id="t1", user_id=1, is_admin=False, log_file=tmp_path / "nope.json") + + +@pytest.mark.asyncio +async def test_get_task_logs_rejects_out_of_range_max_lines(tmp_path): + # The shared core.collect_task_logs enforces the 1..MAX_TASK_LOG_LINES bound, + # same as the admin task-logs route. + log = tmp_path / "app.json" + log.write_text("") + svc = _service(jobs=FakeJobs(details={"user_id": 1})) + with pytest.raises(ValueError): + await svc.get_task_logs(task_id="t1", user_id=1, is_admin=False, log_file=log, max_lines=10_000) + + +# --------------------------------------------------------------------------- +# Chunk lookup +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_chunk_by_id_not_found(): + with pytest.raises(KeyError): + await _service(conversion=FakeConversion(chunk=None)).get_chunk_by_id(chunk_id="9", allowed_partitions=["a"]) + + +@pytest.mark.asyncio +async def test_get_chunk_by_id_partition_denied(): + chunk = {"page_content": "x", "metadata": {"partition": "secret"}} + with pytest.raises(PermissionError): + await _service(conversion=FakeConversion(chunk=chunk)).get_chunk_by_id(chunk_id="9", allowed_partitions=["a"]) + + +@pytest.mark.asyncio +async def test_get_chunk_by_id_allows_in_scope(): + chunk = {"page_content": "x", "metadata": {"partition": "a"}} + out = await _service(conversion=FakeConversion(chunk=chunk)).get_chunk_by_id(chunk_id="9", allowed_partitions=["a"]) + assert out["page_content"] == "x" + + +# --------------------------------------------------------------------------- +# Write operations + editor ACL +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_delete_file_requires_editor(): + parts = FakePartitions(members=[{"user_id": 1, "role": "viewer"}]) + with pytest.raises(PermissionError): + await _service(partitions=parts).delete_file(partition="a", file_id="f1", allowed_partitions=["a"], user_id=1) + + +@pytest.mark.asyncio +async def test_delete_file_editor_succeeds(): + parts = FakePartitions(members=[{"user_id": 1, "role": "editor"}], exists=True) + indexing = FakeIndexing() + out = await _service(partitions=parts, indexing=indexing).delete_file( + partition="a", file_id="f1", allowed_partitions=["a"], user_id=1 + ) + assert indexing.deleted == [("f1", "a")] + assert "deleted" in out["message"] + + +@pytest.mark.asyncio +async def test_delete_file_admin_wildcard_bypasses_membership(): + indexing = FakeIndexing() + await _service(partitions=FakePartitions(exists=True), indexing=indexing).delete_file( + partition="a", file_id="f1", allowed_partitions=["all"], user_id=1 + ) + assert indexing.deleted == [("f1", "a")] + + +@pytest.mark.asyncio +async def test_delete_file_invalid_file_id(): + with pytest.raises(ValueError): + await _service().delete_file(partition="a", file_id="bad/id", allowed_partitions=["all"], user_id=1) + + +@pytest.mark.asyncio +async def test_copy_file_dest_exists_raises(): + parts = FakePartitions(exists=True) # both source and dest "exist" + with pytest.raises(FileExistsError): + await _service(partitions=parts).copy_file( + source_partition="a", + source_file_id="s", + dest_partition="a", + dest_file_id="d", + allowed_partitions=["all"], + user_id=1, + ) + + +@pytest.mark.asyncio +async def test_copy_file_success(): + # source exists, dest does not + class P(FakePartitions): + async def file_exists(self, file_id, partition): + return file_id == "s" + + indexing = FakeIndexing() + out = await _service(partitions=P(), indexing=indexing).copy_file( + source_partition="a", + source_file_id="s", + dest_partition="b", + dest_file_id="d", + allowed_partitions=["all"], + user_id=1, + ) + assert indexing.copied[0]["source_file_id"] == "s" + assert indexing.copied[0]["target_partition"] == "b" + assert out["dest_file_id"] == "d" + + +# --------------------------------------------------------------------------- +# index_url +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_index_url_rejects_non_http_scheme(): + with pytest.raises(ValueError): + await _service(partitions=FakePartitions(exists=False)).index_url( + url="ftp://x/y.pdf", partition="a", file_id="f1", allowed_partitions=["all"], user_id=1 + ) + + +@pytest.mark.asyncio +async def test_index_url_existing_file_raises(monkeypatch): + parts = FakePartitions(exists=True) + with pytest.raises(FileExistsError): + await _service(partitions=parts).index_url( + url="https://x/y.pdf", partition="a", file_id="f1", allowed_partitions=["all"], user_id=1 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "http://127.0.0.1/x.pdf", # loopback + "http://localhost/x.pdf", # localhost name + "http://169.254.169.254/latest/meta-data", # cloud metadata + "http://10.0.0.5/x.pdf", # private + "http://2130706433/x.pdf", # decimal-encoded 127.0.0.1 + "file:///etc/passwd", # non-http scheme + ], +) +async def test_index_url_blocks_ssrf_targets(url): + parts = FakePartitions(exists=False, partition_exists=False) + with pytest.raises(ValueError): + await _service(partitions=parts).index_url( + url=url, partition="newpart", file_id="f1", allowed_partitions=["all"], user_id=1 + ) + + +@pytest.mark.asyncio +async def test_assert_host_resolves_safely_blocks_dns_to_private(monkeypatch): + from services.orchestrators import mcp_service as mcp_mod + + class FakeLoop: + async def getaddrinfo(self, host, port, type=None): + return [(0, 0, 0, "", ("169.254.169.254", 0))] # cloud metadata + + monkeypatch.setattr(mcp_mod.asyncio, "get_running_loop", lambda: FakeLoop()) + with pytest.raises(ValueError): + await MCPService._assert_host_resolves_safely("metadata.evil.example") + + +@pytest.mark.asyncio +async def test_assert_host_resolves_safely_allows_public(monkeypatch): + from services.orchestrators import mcp_service as mcp_mod + + class FakeLoop: + async def getaddrinfo(self, host, port, type=None): + return [(0, 0, 0, "", ("8.8.8.8", 0))] + + monkeypatch.setattr(mcp_mod.asyncio, "get_running_loop", lambda: FakeLoop()) + await MCPService._assert_host_resolves_safely("dns.example") # no raise + + +@pytest.mark.asyncio +async def test_index_url_auto_creates_partition_and_indexes(monkeypatch): + parts = FakePartitions(exists=False, partition_exists=False) + indexing = FakeIndexing() + svc = _service(partitions=parts, indexing=indexing) + + async def fake_download(url, dest): + dest.write_bytes(b"data") + + monkeypatch.setattr(svc, "_safe_download", fake_download) + + out = await svc.index_url( + url="https://example.com/report.pdf", + partition="newpart", + file_id="f1", + allowed_partitions=["other"], + user_id=7, + extra_metadata={"author": "me", "created_by": 999}, # created_by must be stripped + ) + # auto-created the missing partition, owned by the caller + assert parts.created == [("newpart", 7)] + added = indexing.added[0] + assert added["file_id"] == "f1" + assert added["partition"] == "newpart" + assert added["sanitized_filename"] == "report.pdf" + assert added["metadata"]["source_url"] == "https://example.com/report.pdf" + assert added["metadata"]["author"] == "me" + assert "created_by" not in added["metadata"] # protected key dropped + assert out["task_id"] == "task-123" + + +@pytest.mark.asyncio +async def test_safe_download_rejects_redirect_to_private(monkeypatch, tmp_path): + # A public host that 302-redirects to a loopback target must be rejected on + # the re-validated hop, not followed. + from services.orchestrators import mcp_service as mcp_mod + + def handler(request): + return httpx.Response(302, headers={"location": "http://127.0.0.1/secret"}) + + transport = httpx.MockTransport(handler) + real_async_client = mcp_mod.httpx.AsyncClient + + def factory(*args, **kwargs): + kwargs["transport"] = transport + return real_async_client(*args, **kwargs) + + monkeypatch.setattr(mcp_mod.httpx, "AsyncClient", factory) + + svc = _service() + + async def _resolves_ok(host): # skip real DNS for the initial public host + return None + + monkeypatch.setattr(svc, "_assert_host_resolves_safely", _resolves_ok) + + with pytest.raises(ValueError, match="disallowed"): + await svc._safe_download("http://public.example/file.pdf", tmp_path / "out") + + +# --------------------------------------------------------------------------- +# Metadata hardening + input bounds +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_update_metadata_strips_protected_keys_keeps_move(): + indexing = FakeIndexing() + await _service(partitions=FakePartitions(exists=True), indexing=indexing).update_file_metadata( + partition="a", + file_id="f1", + metadata={"author": "x", "source": "/evil", "created_by": 999, "partition": "dest"}, + allowed_partitions=["all"], + user_id=1, + ) + _file_id, sent_md, _partition, _user = indexing.updated[0] + assert sent_md["author"] == "x" + assert sent_md["partition"] == "dest" # authorized move control preserved + assert "source" not in sent_md and "created_by" not in sent_md + + +@pytest.mark.asyncio +async def test_update_metadata_move_requires_editor_on_destination(): + # Non-wildcard caller: editor on source but only viewer on the move target. + class P(FakePartitions): + async def list_members(self, partition): + role = "viewer" if partition == "dest" else "editor" + return [{"user_id": 5, "role": role}] + + svc = _service(partitions=P(exists=True)) + with pytest.raises(PermissionError): + await svc.update_file_metadata( + partition="src", + file_id="f1", + metadata={"partition": "dest"}, + allowed_partitions=["src", "dest"], + user_id=5, + ) + + +@pytest.mark.asyncio +async def test_get_file_chunks_rejects_bad_offset_and_limit(): + svc = _service(partitions=FakePartitions(exists=True)) + with pytest.raises(ValueError): + await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], offset=-1) + with pytest.raises(ValueError): + await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], limit=0) + + +@pytest.mark.asyncio +async def test_get_file_chunks_caps_page_size(): + from services.orchestrators.mcp_service import _MAX_CHUNKS_PER_CALL + + rows = [{"_id": i, "text": f"c{i}"} for i in range(_MAX_CHUNKS_PER_CALL + 50)] + svc = _service(vector_store=FakeVectorStore(rows=rows)) + out = await svc.get_file_chunks(partition="a", file_id="f1", allowed_partitions=["a"], offset=0, limit=-1) + assert len(out["chunks"]) == _MAX_CHUNKS_PER_CALL + assert out["has_more"] is True diff --git a/openrag/services/storage/milvus_store.py b/openrag/services/storage/milvus_store.py index 58f89a4bd..ab8139845 100644 --- a/openrag/services/storage/milvus_store.py +++ b/openrag/services/storage/milvus_store.py @@ -683,7 +683,14 @@ def _parse_search_response(self, response: Any) -> list[dict[str, Any]]: for hit in response[0]: entity = hit.get("entity", {}) if isinstance(hit, dict) else {} record = {k: v for k, v in entity.items() if k not in _SEARCH_RESULT_DROPPED_KEYS} - record["id"] = self._milvus_id_to_str(hit.get("id")) + # The primary key field is named ``_id`` (auto_id INT64), so Milvus + # exposes it on the hit as ``_id`` and also inside ``entity`` — NOT + # under the generic ``id`` key. Reading ``id`` yielded a literal + # "None" string for every search result. + pk = hit.get("_id") + if pk is None: + pk = entity.get("_id") + record["id"] = self._milvus_id_to_str(pk) if pk is not None else None record["score"] = hit.get("distance") out.append(record) return out diff --git a/openrag/services/storage/test_milvus_store.py b/openrag/services/storage/test_milvus_store.py index c66e48258..1908cfd36 100644 --- a/openrag/services/storage/test_milvus_store.py +++ b/openrag/services/storage/test_milvus_store.py @@ -444,3 +444,37 @@ async def test_hybrid_store_requires_query_text(self, store: MilvusVectorStore) """ with pytest.raises(VDBSearchError, match="query_text"): await store.search([0.1, 0.2], collection="default") + + +# --------------------------------------------------------------------------- +# _parse_search_response +# --------------------------------------------------------------------------- + + +class TestParseSearchResponse: + """Milvus 2.6 exposes the ``_id`` auto-id PK on the hit (and in the entity), + never under the generic ``id`` key — the parser must surface the real id.""" + + def test_id_taken_from_hit_underscore_id(self, store: MilvusVectorStore) -> None: + hit = { + "_id": 466609479666371445, + "distance": 0.42, + "entity": {"_id": 466609479666371445, "text": "hello", "file_id": "doc1", "vector": [0.1, 0.2]}, + } + (record,) = store._parse_search_response([[hit]]) + assert record["id"] == "466609479666371445" # not the literal "None" + assert record["score"] == 0.42 + assert "vector" not in record + + def test_id_falls_back_to_entity_underscore_id(self, store: MilvusVectorStore) -> None: + hit = {"distance": 0.1, "entity": {"_id": 123, "text": "t"}} + (record,) = store._parse_search_response([[hit]]) + assert record["id"] == "123" + + def test_missing_pk_yields_none_not_the_string(self, store: MilvusVectorStore) -> None: + hit = {"distance": 0.1, "entity": {"text": "t"}} + (record,) = store._parse_search_response([[hit]]) + assert record["id"] is None + + def test_empty_response_is_empty_list(self, store: MilvusVectorStore) -> None: + assert store._parse_search_response([]) == [] diff --git a/openrag/utils/logger.py b/openrag/utils/logger.py index 64771b9e8..a43031636 100644 --- a/openrag/utils/logger.py +++ b/openrag/utils/logger.py @@ -2,6 +2,7 @@ import sys from config import load_config +from core.utils.log_tail import app_log_file from loguru import logger @@ -44,11 +45,11 @@ def formatter(record): logger.add(sys.stderr, format=formatter, level=config.verbose.level, colorize=False) # JSON logs to file for later use (e.g. Grafana ingestion) - log_dir = config.paths.log_dir if hasattr(config.paths, "log_dir") else "logs" + log_path = app_log_file(getattr(config.paths, "log_dir", None)) try: - os.makedirs(log_dir, exist_ok=True) + os.makedirs(log_path.parent, exist_ok=True) logger.add( - f"{log_dir}/app.json", + str(log_path), serialize=True, level=config.verbose.level, rotation="10 MB", diff --git a/pyproject.toml b/pyproject.toml index caa14330a..78488c3ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ dependencies = [ "cryptography>=42", "tenacity>=8.2.0", "aiobreaker>=1.2.0", + "mcp>=1.11.0", ] [dependency-groups] diff --git a/uv.lock b/uv.lock index 3f088049d..dbb146d56 100644 --- a/uv.lock +++ b/uv.lock @@ -2674,6 +2674,7 @@ dependencies = [ { name = "lxml" }, { name = "marker-pdf" }, { name = "markitdown", extra = ["docx"] }, + { name = "mcp" }, { name = "numba" }, { name = "openai" }, { name = "pip" }, @@ -2740,6 +2741,7 @@ requires-dist = [ { name = "lxml", specifier = ">=5.0.0" }, { name = "marker-pdf", specifier = ">=0.2.17" }, { name = "markitdown", extras = ["docx"], specifier = ">=0.1.3" }, + { name = "mcp", specifier = ">=1.11.0" }, { name = "numba", specifier = ">=0.61.2" }, { name = "openai", specifier = ">=1.64.0" }, { name = "pip", specifier = ">=25.0.1" },