diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0b4671aae..4daddb2de 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -22,12 +22,20 @@ jobs: python-version: ["3.10", "3.13"] agent-frameworks: [without, with] pdfium: ["5"] + mcp: [latest] include: # one leg holds the 4.x insurance line - python-version: "3.10" agent-frameworks: with pdfium: "4" - name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks, pdfium ${{ matrix.pdfium }}) + mcp: latest + # one leg holds the declared mcp floor: the local MCP server's + # 1.x registration path never runs on the latest release + - python-version: "3.10" + agent-frameworks: without + pdfium: "5" + mcp: "1.19.0" + name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks, pdfium ${{ matrix.pdfium }}, mcp ${{ matrix.mcp }}) timeout-minutes: 15 steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 @@ -36,6 +44,7 @@ jobs: python-version: ${{ matrix.python-version }} cache: pip - run: pip install -r requirements.txt pytest + - run: pip install --no-deps -e . - if: matrix.agent-frameworks == 'without' # requirements.txt carries it; this leg tests the no-framework paths run: pip uninstall -y openai-agents @@ -43,6 +52,8 @@ jobs: run: pip install openai-agents claude-agent-sdk anthropic - if: matrix.pdfium == '4' run: pip install "pypdfium2<5" + - if: matrix.mcp != 'latest' + run: pip install "mcp==${{ matrix.mcp }}" - run: python -m pytest -q env: PAGEINDEX_API_KEY: ${{ secrets.PAGEINDEX_API_KEY }} diff --git a/README.md b/README.md index 05457ee57..f9545444a 100644 --- a/README.md +++ b/README.md @@ -109,6 +109,29 @@ Configure other models, streaming, multi-document search, citations, and more. Drop PageIndex tools into the OpenAI Agents SDK, the Claude Agent SDK, or any other framework. +### Local MCP server + +Serve a local document store to Claude Desktop, Cursor, or any other MCP client over stdio: + +```bash +pageindex-mcp --storage-path /absolute/path/to/index-store +``` + +`--storage-path` is the `storage_path` your documents were indexed into with `PageIndexLocalClient`. Serving needs no model API key: the MCP client's own model calls the tools. The tools are read-only by default (`browse_documents`, `get_document`, `get_document_structure`, `get_page_content`); add `--management` to enable `remove_document`. + +```json +{ + "mcpServers": { + "pageindex-local": { + "command": "/absolute/path/to/.venv/bin/pageindex-mcp", + "args": ["--storage-path", "/absolute/path/to/index-store"] + } + } +} +``` + +Use absolute paths: desktop apps start the server without your shell's working directory or `PATH`. + # Benchmarks @@ -205,7 +228,7 @@ print(client.chat("What was the 2023 operating margin?", doc_id=doc_id)) | OCR & image understanding | — | ✓ | | [Metadata](https://docs.pageindex.ai/sdk/documents#metadata-cloud) | — | ✓ | | [Folders](https://docs.pageindex.ai/sdk/documents#folders-cloud) | — | ✓ | -| [MCP server](https://docs.pageindex.ai/mcp) | — | ✓ | +| [MCP server](https://docs.pageindex.ai/mcp) | [stdio](#local-mcp-server) | ✓ | ### More About PageIndex Cloud diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 26128127c..ed8e5decd 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -31,7 +31,8 @@ } _SUBMODULES = {"agent_tools", "chat_stream", "client", "cloud_api", "errors", "flash", "imaging", "integrations", "local_api", "local_chat", - "local_store", "mcp_bridge", "page_index_classic", + "local_mcp_server", "local_store", "mcp_bridge", + "page_index_classic", "page_index_md", "tree_optimize", "types", "utils"} diff --git a/pageindex/local_mcp_server.py b/pageindex/local_mcp_server.py new file mode 100644 index 000000000..9ac9d6c25 --- /dev/null +++ b/pageindex/local_mcp_server.py @@ -0,0 +1,168 @@ +"""MCP tools over an existing local PageIndex document store. + +The server reuses the SDK's tool contract and dispatchers. It does not index +documents or run a chat model; an MCP host uses the tools to read documents +that were previously indexed with a PageIndexLocalClient. +""" + +import asyncio +import io +import os +import sys + +import anyio +from mcp import types +from mcp.server.lowlevel import Server + +from ._version import sdk_version +from .agent_tools import TOOL_CONTRACT, _tool_specs + + +class LocalMcpServer(Server): + """Serve local document tools through MCP 1.x or 2.x. + + Args: + client: A PageIndexLocalClient connected to the indexed document store. + include_management: Expose and allow remove_document when True. + The default tool set is read-only. + + Constructing this object does not start a transport: await + ``serve_stdio()``, or run the ``pageindex-mcp`` command. + """ + + def __init__(self, client, include_management: bool = False): + from .client import PageIndexLocalClient + + if not isinstance(client, PageIndexLocalClient): + raise TypeError("LocalMcpServer requires a PageIndexLocalClient") + self.specs = _tool_specs(client, include_management) + self.invokers = { + name: invoke + for name, _, _, invoke in self.specs + } + + options = { + "version": sdk_version(), + "instructions": client.agent_instructions( + include_management=include_management), + } + # MCP 1.x registers decorators on an instance; 2.x accepts callbacks. + if hasattr(Server, "list_tools"): + super().__init__("pageindex-local-mcp", **options) + + async def list_handler(): + return await self.list_tools(None, None) + + async def call_handler(name, arguments): + params = types.CallToolRequestParams(name=name, arguments=arguments) + return await self.call_tool(None, params) + + Server.list_tools(self)(list_handler) + Server.call_tool(self, validate_input=False)(call_handler) + else: + super().__init__( + "pageindex-local-mcp", **options, + on_list_tools=self.list_tools, + on_call_tool=self.call_tool, + ) + + async def list_tools(self, context, params): + """Return registered local schemas and their MCP safety annotations. + + ``context`` and pagination ``params`` are supplied by MCP 2.x. The + fixed local tool catalog fits in one response, so neither is needed. + """ + return types.ListToolsResult(tools=[ + types.Tool( + name=name, + description=description, + inputSchema=schema, + annotations=types.ToolAnnotations( + **TOOL_CONTRACT[name].get("annotations", {})), + ) + for name, description, schema, _ in self.specs + ]) + + async def call_tool(self, context, params): + """Dispatch a registered tool off the event loop and preserve errors. + + ``params`` carries the tool name and arguments. Only registered + invokers can execute, so a management tool remains disabled even if + a client calls its name directly. ``context`` is unused. + """ + invoke = self.invokers.get(params.name) + if invoke is None: + return types.CallToolResult( + content=[types.TextContent( + type="text", text=f"Unknown or disabled tool: {params.name}")], + isError=True + ) + + blocks, is_error = await asyncio.to_thread(invoke, params.arguments or {}) + return types.CallToolResult.model_validate( + {"content": blocks, "isError": is_error}) + + async def serve_stdio(self): + """Serve over stdin/stdout until the host closes the pipe. + + While serving, fd 1 points at stderr and the protocol writes to a + private duplicate of the original stdout, so stray output from tools, + C extensions or child processes never corrupts the JSON-RPC stream. + MCP 2.x's stdio_server diverts stdout itself; 1.x does not. + """ + from mcp.server.stdio import stdio_server + + sys.stdout.flush() + protocol_fd = os.dup(1) + os.dup2(2, 1) + protocol_out = io.TextIOWrapper( + os.fdopen(protocol_fd, "wb", closefd=False), encoding="utf-8") + try: + async with stdio_server(stdout=anyio.wrap_file(protocol_out)) as ( + read_stream, write_stream): + await self.run(read_stream, write_stream, + self.create_initialization_options()) + finally: + sys.stdout.flush() + try: + protocol_out.flush() + except (OSError, ValueError): + pass # host already closed the pipe + os.dup2(protocol_fd, 1) + os.close(protocol_fd) + + +def main(): + """Launch the local stdio server from the ``pageindex-mcp`` command.""" + from argparse import ArgumentParser + + from .client import PageIndexLocalClient + from .errors import PageIndexAPIError + + parser = ArgumentParser(description="Serve a local PageIndex document store over MCP stdio.") + parser.add_argument("--storage-path", required=True, + help="Path to an existing indexed document store.") + parser.add_argument("--management", action="store_true", + help="Enable document deletion (disabled by default).") + args = parser.parse_args() + # The store reads as an empty library when it is missing or unreadable, + # so a bad path would otherwise serve nothing without complaint. + if not os.path.isdir(args.storage_path): + parser.error(f"--storage-path {args.storage_path!r} is not a directory") + if not os.access(args.storage_path, os.R_OK | os.X_OK): + parser.error(f"--storage-path {args.storage_path!r} is not readable") + if not os.path.isfile(os.path.join(args.storage_path, "manifest.json")): + # Not fatal: a fresh store has no manifest until the first document. + print(f"pageindex-mcp: warning: no indexed documents in " + f"{args.storage_path!r}; index with PageIndexLocalClient first", + file=sys.stderr) + + try: + client = PageIndexLocalClient(storage_path=args.storage_path) + mcp_server = LocalMcpServer(client, include_management=args.management) + except (PageIndexAPIError, OSError) as exc: + parser.exit(1, f"pageindex-mcp: error: {exc}\n") + try: + asyncio.run(mcp_server.serve_stdio()) + except KeyboardInterrupt: + sys.exit(130) diff --git a/pyproject.toml b/pyproject.toml index ef21a2ea7..c4f76fa13 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,9 @@ Homepage = "https://pageindex.ai" Documentation = "https://docs.pageindex.ai" Issues = "https://github.com/VectifyAI/PageIndex/issues" +[tool.poetry.scripts] +pageindex-mcp = "pageindex.local_mcp_server:main" + [build-system] requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" diff --git a/tests/test_local_mcp.py b/tests/test_local_mcp.py new file mode 100644 index 000000000..8f8291601 --- /dev/null +++ b/tests/test_local_mcp.py @@ -0,0 +1,316 @@ +"""Local MCP protocol tests against a seeded store, without LLM calls.""" +import asyncio +import json +import os +import shutil +import subprocess +import sys +import sysconfig +from pathlib import Path + +import anyio +import pytest +from mcp import ClientSession, StdioServerParameters, types +from mcp.client.stdio import stdio_client +from mcp.shared.memory import create_client_server_memory_streams +from test_agent_tools import seed_doc + +from pageindex import PageIndexCloudClient, PageIndexLocalClient +from pageindex.agent_tools import _local_schema, tool_names +from pageindex.local_mcp_server import LocalMcpServer + + +@pytest.fixture +def local_client(tmp_path): + storage = str(tmp_path / "store") + seed_doc(storage, "pi-a", "report.pdf") + return PageIndexLocalClient(storage_path=storage) + + +def wire(result): + """MCP 1.x and 2.x use different Python names but the same wire aliases.""" + return result.model_dump(by_alias=True, exclude_none=True) + + +def payload(result): + return json.loads(wire(result)["content"][0]["text"]) + + +@pytest.fixture +def executable(): + """Use the installed entry point, including its generated console wrapper.""" + suffix = ".exe" if sys.platform == "win32" else "" + command = Path(sysconfig.get_path("scripts")) / f"pageindex-mcp{suffix}" + if command.is_file(): + return str(command) + found = shutil.which("pageindex-mcp") + if found: + return found + # CI installs the checkout, so a missing command there is a real failure. + if os.environ.get("CI"): + pytest.fail("pageindex-mcp is not installed") + pytest.skip("pageindex-mcp is not installed; run python -m pip install -e .") + + +@pytest.mark.parametrize("management", [False, True]) +def test_installed_executable_round_trip(executable, local_client, tmp_path, management): + """The command selects the requested store and gates deletion over stdio.""" + async def check(): + with anyio.fail_after(15): + args = ["--storage-path", str(local_client.storage_path)] + if management: + args.append("--management") + parameters = StdioServerParameters( + command=executable, args=args, cwd=str(tmp_path), + ) + async with stdio_client(parameters) as streams, ClientSession(*streams) as session: + await session.initialize() + catalog = await session.list_tools() + assert [tool.name for tool in catalog.tools] == list(tool_names(management)) + result = await session.call_tool("browse_documents", {}) + assert not wire(result)["isError"] + assert payload(result)["documents"][0]["name"] == "report.pdf" + removed = await session.call_tool( + "remove_document", {"doc_names": ["report.pdf"]}, + ) + assert wire(removed)["isError"] is (not management) + remaining = await session.call_tool("browse_documents", {}) + assert len(payload(remaining)["documents"]) == (0 if management else 1) + + asyncio.run(check()) + + +def test_executable_help_and_invalid_arguments(executable, tmp_path): + help_result = subprocess.run( + [executable, "--help"], capture_output=True, text=True, timeout=10, + ) + assert help_result.returncode == 0 + assert "--storage-path" in help_result.stdout + assert "--management" in help_result.stdout + for arguments in [ + [], + ["--storage-path", str(tmp_path), "--management", "false"], + ["--storage-path", str(tmp_path / "missing")], + ]: + result = subprocess.run( + [executable, *arguments], capture_output=True, text=True, timeout=10, + ) + assert result.returncode == 2 + assert result.stdout == "" + assert "error:" in result.stderr + + +@pytest.mark.skipif(sys.platform == "win32" or os.geteuid() == 0, + reason="needs POSIX permissions and a non-root user") +def test_executable_rejects_unreadable_store(executable, tmp_path): + store = tmp_path / "locked" + store.mkdir() + store.chmod(0) + try: + result = subprocess.run( + [executable, "--storage-path", str(store)], + capture_output=True, text=True, timeout=10, + ) + finally: + store.chmod(0o755) + assert result.returncode == 2 + assert result.stdout == "" + assert "is not readable" in result.stderr + + +def test_executable_warns_on_empty_store_and_exits_on_eof(executable, tmp_path): + """An empty store still serves, but says so on stderr; closing stdin, as + a host does, shuts the server down cleanly.""" + result = subprocess.run( + [executable, "--storage-path", str(tmp_path)], + input="", capture_output=True, text=True, timeout=15, + ) + assert result.returncode == 0 + assert result.stdout == "" + assert "warning: no indexed documents" in result.stderr + + +async def round_trip(server, assertions): + """Exercise real initialization and dispatch with a bounded lifetime.""" + with anyio.fail_after(10): + async with create_client_server_memory_streams() as (client_streams, server_streams), \ + anyio.create_task_group() as tasks: + tasks.start_soon( + server.run, *server_streams, + server.create_initialization_options(), + ) + async with ClientSession(*client_streams) as session: + initialized = await session.initialize() + await assertions(session, initialized) + tasks.cancel_scope.cancel() + + +def test_protocol_discovery_and_document_reads(local_client): + async def assertions(session, initialized): + assert initialized.instructions == local_client.agent_instructions() + assert initialized.capabilities.tools is not None + catalog = await session.list_tools() + assert [tool.name for tool in catalog.tools] == list(tool_names()) + for tool in catalog.tools: + definition = wire(tool) + assert definition["inputSchema"] == _local_schema(tool.name) + assert definition["annotations"]["readOnlyHint"] is True + + browsed = await session.call_tool("browse_documents", {}) + assert not wire(browsed)["isError"] + assert payload(browsed)["documents"][0]["name"] == "report.pdf" + + for name, arguments in [ + ("get_document", {"doc_name": "report.pdf"}), + ("get_document_structure", {"doc_name": "report.pdf"}), + ("get_page_content", {"doc_name": "report.pdf", "pages": "1-2"}), + ]: + result = await session.call_tool(name, arguments) + assert not wire(result)["isError"] + assert payload(result)["success"] is True + if name == "get_page_content": + assert "apples" in wire(result)["content"][0]["text"] + assert "bananas" in wire(result)["content"][0]["text"] + + asyncio.run(round_trip(LocalMcpServer(local_client), assertions)) + + +def test_protocol_errors_and_management_gate(local_client): + async def assertions(session, initialized): + missing = await session.call_tool("get_document", {"doc_name": "missing.pdf"}) + assert wire(missing)["isError"] is True + assert payload(missing)["errorCode"] == "NOT_FOUND" + + for name, arguments in [ + ("unknown_tool", {}), + ("remove_document", {"doc_names": ["report.pdf"]}), + ("get_page_content", {}), + ]: + result = await session.call_tool(name, arguments) + assert wire(result)["isError"] is True + + # Calling a hidden management tool must not delete anything. + assert local_client.get_document("pi-a")["name"] == "report.pdf" + + asyncio.run(round_trip(LocalMcpServer(local_client), assertions)) + + +def test_management_opt_in_allows_deletion(local_client): + async def assertions(session, initialized): + catalog = await session.list_tools() + assert [tool.name for tool in catalog.tools] == list(tool_names(True)) + result = await session.call_tool("remove_document", {"doc_names": ["report.pdf"]}) + assert not wire(result)["isError"] + assert payload(result)["success"] is True + browsed = await session.call_tool("browse_documents", {}) + assert payload(browsed)["documents"] == [] + + asyncio.run(round_trip(LocalMcpServer(local_client, True), assertions)) + + +def test_string_booleans_reach_the_tool_layer(local_client): + """Models often send booleans as strings; the tool layer coerces them, as + on cloud. SDK-side schema validation must not reject them first.""" + async def assertions(session, initialized): + result = await session.call_tool( + "get_document", + {"doc_name": "report.pdf", "wait_for_completion": "false"}, + ) + assert not wire(result)["isError"], wire(result)["content"][0]["text"] + assert payload(result)["success"] is True + + asyncio.run(round_trip(LocalMcpServer(local_client), assertions)) + + +def test_cloud_client_is_rejected(): + with pytest.raises(TypeError, match="PageIndexLocalClient"): + LocalMcpServer(PageIndexCloudClient(api_key="test-key")) + + +def test_invoker_runs_off_event_loop(local_client): + import threading + + server = LocalMcpServer(local_client) + main_thread = threading.get_ident() + worker_threads = [] + + def invoke(arguments): + worker_threads.append(threading.get_ident()) + return [{"type": "text", "text": "ok"}], False + + server.invokers["browse_documents"] = invoke + result = asyncio.run(server.call_tool( + None, types.CallToolRequestParams(name="browse_documents", arguments={}), + )) + assert wire(result)["content"][0]["text"] == "ok" + assert len(worker_threads) == 1 and worker_threads[0] != main_thread + + +def test_stdio_subprocess_round_trip(local_client): + """Catch stdout contamination and startup/shutdown issues over real pipes.""" + source = """ +import asyncio +import sys +from pageindex import PageIndexLocalClient +from pageindex.local_mcp_server import LocalMcpServer + +async def main(): + server = LocalMcpServer(PageIndexLocalClient(storage_path=sys.argv[1])) + await server.serve_stdio() + +asyncio.run(main()) +""" + + async def check(): + with anyio.fail_after(15): + parameters = StdioServerParameters( + command=sys.executable, + args=["-c", source, str(local_client.storage_path)], + ) + async with stdio_client(parameters) as streams, \ + ClientSession(*streams) as session: + await session.initialize() + assert [tool.name for tool in (await session.list_tools()).tools] == list(tool_names()) + result = await session.call_tool("browse_documents", {}) + assert payload(result)["documents"][0]["name"] == "report.pdf" + + asyncio.run(check()) + + +def test_stdio_survives_stray_stdout_from_a_tool(local_client): + """A tool that writes to stdout mid-call must not corrupt the JSON-RPC + stream; no trailing newline, so the noise would fuse with the reply.""" + source = """ +import asyncio +import sys +from pageindex import PageIndexLocalClient +from pageindex.local_mcp_server import LocalMcpServer + +async def main(): + server = LocalMcpServer(PageIndexLocalClient(storage_path=sys.argv[1])) + browse = server.invokers["browse_documents"] + + def noisy(arguments): + print("stray tool output", end="", flush=True) + return browse(arguments) + + server.invokers["browse_documents"] = noisy + await server.serve_stdio() + +asyncio.run(main()) +""" + + async def check(): + with anyio.fail_after(15): + parameters = StdioServerParameters( + command=sys.executable, + args=["-c", source, str(local_client.storage_path)], + ) + async with stdio_client(parameters) as streams, \ + ClientSession(*streams) as session: + await session.initialize() + for _ in range(2): # the second reply would carry the noise + result = await session.call_tool("browse_documents", {}) + assert payload(result)["documents"][0]["name"] == "report.pdf" + + asyncio.run(check())