diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 173ed9fd8..81b128752 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -183,7 +183,7 @@ jobs: - name: Build Ultraplot run: | - pip install --no-build-isolation --no-deps . + pip install --no-build-isolation ".[mcp]" - name: Run full coverage suite run: | diff --git a/README.rst b/README.rst index 4ca53f96c..2a6ef838d 100644 --- a/README.rst +++ b/README.rst @@ -112,7 +112,7 @@ pyCirclize-based plots require the optional ``circos`` extra: The ``docs`` extra also includes pyCirclize for building the documentation. -To install all optional dependency groups (``circos``, ``docs``, and ``stats``): +To install the ``circos``, ``docs``, and ``stats`` dependency groups together: .. code-block:: bash @@ -131,6 +131,113 @@ To install a development version of UltraPlot, you can use or clone the repository and run ``pip install -e .`` inside the ``ultraplot`` folder. +MCP server +========== + +UltraPlot includes a Model Context Protocol (MCP) server that lets AI assistants +search documentation and examples, inspect the live Python API, and read source +code and release notes. + +Run directly with uvx +--------------------- + +With `uv `__ installed, +your MCP client can launch the server with ``uvx``. uv installs the package and +its dependencies in an isolated environment automatically, so you do not need +to create a virtual environment or install UltraPlot separately. + +For a PyPI release containing the MCP server, the launch command is: + +.. code-block:: bash + + uvx --from 'ultraplot[mcp]' ultraplot-mcp + +For clients that use an ``mcpServers`` configuration, add: + +.. code-block:: json + + { + "mcpServers": { + "ultraplot": { + "command": "uvx", + "args": ["--from", "ultraplot[mcp]", "ultraplot-mcp"] + } + } + } + +The client starts the server when needed and communicates with it over stdio. +Other clients may use a different configuration format; use the same command +and arguments. ``uvx`` is equivalent to ``uv tool run``. + +Until the MCP server is released on PyPI, run it directly from the feature +branch instead: + +.. code-block:: bash + + uvx --from 'ultraplot[mcp] @ git+https://github.com/ultraplot/ultraplot.git@feat/mcp' ultraplot-mcp + +For this development version, replace ``ultraplot[mcp]`` in the client +configuration with +``ultraplot[mcp] @ git+https://github.com/ultraplot/ultraplot.git@feat/mcp``. + +Install persistently with uv +---------------------------- + +Alternatively, keep the executable on your ``PATH`` by installing it as a uv +tool. For a PyPI release containing the MCP server: + +.. code-block:: bash + + uv tool install 'ultraplot[mcp]' + ultraplot-mcp --help + +Before that release, install from the feature branch: + +.. code-block:: bash + + uv tool install 'ultraplot[mcp] @ git+https://github.com/ultraplot/ultraplot.git@feat/mcp' + +Then configure your client to launch ``ultraplot-mcp`` with no arguments. +If uv reports that its executable directory is missing from ``PATH``, run +``uv tool update-shell`` and restart your shell. + +Install from a checkout +----------------------- + +From a checkout containing the MCP implementation, install the optional ``mcp`` +extra in the Python environment you want the server to use: + +.. code-block:: bash + + pip install -e '.[mcp]' + +Connect an MCP client +--------------------- + +After installing persistently with uv or pip, register the server with an +installed Codex CLI: + +.. code-block:: bash + + ultraplot-mcp install codex + +Restart Codex after registration. Try asking it to search the UltraPlot +examples for shared colorbars or inspect ``ultraplot.subplots``. + +For other MCP clients, configure a stdio server with ``ultraplot-mcp`` as the +command and no arguments. Use the executable's absolute path if the client does +not inherit your Python environment's ``PATH``. Running ``ultraplot-mcp`` starts +the server; ``ultraplot-mcp --help`` lists the available commands. + +Documentation tools read the checkout's ``docs`` directory. Direct uvx and uv +tool installations require a separate documentation checkout for these tools. +If documentation lives elsewhere, set ``ULTRAPLOT_MCP_DOCS`` to its absolute path in the MCP +client's server environment. Documentation is not currently bundled in the +Python package; API and source inspection use the installed UltraPlot version. + +Citing UltraPlot +================ + If you use UltraPlot in your research, please cite the latest release metadata in ``CITATION.cff``. GitHub can export this metadata as BibTeX from the repository's "Cite this repository" panel, and the Zenodo badge below points to diff --git a/pyproject.toml b/pyproject.toml index 7653d3bff..61c66ad42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,9 @@ dynamic = ["version"] "Issue Tracker" = "https://github.com/ultraplot/ultraplot/issues" "Source Code" = "https://github.com/ultraplot/ultraplot" +[project.scripts] +ultraplot-mcp = "ultraplot.mcp:main" + [project.optional-dependencies] all = [ "ultraplot[circos,docs,stats]", @@ -65,6 +68,9 @@ docs = [ "sphinx-sitemap", "typing-extensions" ] +mcp = [ + "mcp>=2.1,<3", +] stats = [ "scipy", ] diff --git a/ultraplot/_lazy.py b/ultraplot/_lazy.py index 7116b9f6a..d1a5044e8 100644 --- a/ultraplot/_lazy.py +++ b/ultraplot/_lazy.py @@ -97,7 +97,8 @@ def _discover_modules(self, module_globals: MutableMapping[str, Any]) -> None: protected.add("figure") for path in base.glob("*.py"): - if path.name.startswith("_") or path.name == "setup.py": + # The MCP CLI is an optional integration, not part of the plotting API. + if path.name.startswith("_") or path.name in {"setup.py", "mcp.py"}: continue module_name = path.stem if module_name in protected: diff --git a/ultraplot/mcp.py b/ultraplot/mcp.py new file mode 100644 index 000000000..17b53b1d7 --- /dev/null +++ b/ultraplot/mcp.py @@ -0,0 +1,695 @@ +from __future__ import annotations + +import inspect +import logging +import os +import pydoc +import re +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +from mcp.server import MCPServer + +# stdout is reserved for MCP JSON-RPC when using the stdio transport. +logging.basicConfig( + level=logging.INFO, + stream=sys.stderr, + format="%(asctime)s %(levelname)s %(message)s", +) + +log = logging.getLogger("ultraplot-mcp") + + +PACKAGE_ROOT = Path(__file__).resolve().parent + +# Development checkout: +# +# ultraplot/ +# ├── docs/ +# ├── ultraplot/ +# │ └── mcp.py +# └── pyproject.toml +# +REPO = Path( + os.environ.get( + "ULTRAPLOT_REPO", + PACKAGE_ROOT.parent, + ) +).resolve() + +SOURCE_DOCS = REPO / "docs" + +# For distributions, the docs can eventually be bundled here: +# +# ultraplot/ +# ├── mcp.py +# └── _mcp_docs/ +# +PACKAGED_DOCS = PACKAGE_ROOT / "_mcp_docs" + +if "ULTRAPLOT_MCP_DOCS" in os.environ: + DOCS = Path(os.environ["ULTRAPLOT_MCP_DOCS"]).resolve() +elif SOURCE_DOCS.is_dir(): + DOCS = SOURCE_DOCS +elif PACKAGED_DOCS.is_dir(): + DOCS = PACKAGED_DOCS +else: + DOCS = SOURCE_DOCS + +log.debug("cwd = %s", Path.cwd()) +log.debug("package root = %s", PACKAGE_ROOT) +log.debug("repo = %s", REPO) +log.debug("docs = %s", DOCS) +log.debug("docs exists = %s", DOCS.exists()) + + +mcp = MCPServer( + "UltraPlot", + instructions=""" +Tools for understanding and using the UltraPlot Python plotting library. + +Prefer UltraPlot-native idioms over equivalent low-level Matplotlib code. + +When answering UltraPlot-specific questions: +1. Inspect the live UltraPlot API when relevant. +2. Search the UltraPlot documentation and examples. +3. Prefer documented behavior over guessing. +4. Inspect source code when documentation is insufficient. + +The tools operate against the UltraPlot version installed alongside this MCP. +""".strip(), +) + + +IGNORED_DOCS = { + "whats_new.rst", + "changelog.rst", + "changes.rst", +} + +STOPWORDS = { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "by", + "different", + "do", + "does", + "for", + "from", + "how", + "i", + "in", + "is", + "it", + "of", + "on", + "or", + "per", + "the", + "to", + "use", + "using", + "what", + "when", + "where", + "which", + "with", +} + + +def _query_terms(query: str) -> list[str]: + """Normalize a search query into meaningful terms.""" + terms = re.findall(r"\w+", query.lower()) + + filtered = [term for term in terms if term not in STOPWORDS and len(term) > 1] + + return filtered or terms + + +def _text_files(): + """Yield searchable documentation files.""" + if not DOCS.is_dir(): + return + + seen: set[Path] = set() + + for pattern in ("**/*.rst", "**/*.md", "**/*.py"): + for path in DOCS.glob(pattern): + if path in seen: + continue + + seen.add(path) + + if path.name.lower() in IGNORED_DOCS: + continue + + yield path + + +def _score_document( + text: str, + query: str, + terms: list[str], + path: Path, +) -> float: + """Calculate a simple relevance score.""" + lower = text.lower() + query_lower = query.lower().strip() + stem = path.stem.lower() + + score = 0.0 + matched_terms = 0 + + # Exact phrase matches are particularly useful. + if query_lower: + score += lower.count(query_lower) * 50 + + for term in terms: + matches = re.findall( + rf"\b{re.escape(term)}\w*\b", + lower, + ) + + occurrences = len(matches) + + if occurrences: + matched_terms += 1 + score += occurrences + + # A query matching the document name is strong evidence. + if term in stem: + score += 20 + + # Reward matching multiple distinct concepts. + score += matched_terms * 5 + + if terms and matched_terms == len(terms): + score += 20 + + return score + + +def _best_match_position( + text: str, + query: str, + terms: list[str], +) -> int: + """Find a useful position for the returned snippet.""" + lower = text.lower() + query_lower = query.lower().strip() + + # Exact phrase first. + if query_lower: + position = lower.find(query_lower) + if position >= 0: + return position + + # Find regions containing several query terms close together. + positions: list[int] = [] + + for term in terms: + for match in re.finditer( + rf"\b{re.escape(term)}\w*\b", + lower, + ): + positions.append(match.start()) + + if not positions: + return 0 + + positions.sort() + + # Find the densest window. + best_position = positions[0] + best_count = 0 + window = 2000 + + for position in positions: + count = sum(position <= other <= position + window for other in positions) + + if count > best_count: + best_count = count + best_position = position + + return best_position + + +def _snippet( + text: str, + query: str, + terms: list[str], + *, + before: int = 600, + after: int = 2400, +) -> str: + """Extract a relevant chunk from a documentation file.""" + position = _best_match_position( + text, + query, + terms, + ) + + start = max(0, position - before) + end = min(len(text), position + after) + + result = text[start:end] + + if start: + result = "...\n" + result + + if end < len(text): + result += "\n..." + + return result + + +def _display_path(path: Path) -> str: + """Return a stable user-facing documentation path.""" + relative = path.relative_to(DOCS) + return str(Path("docs") / relative) + + +def _resolve_doc_path(path: str) -> Path: + """Resolve a user-facing documentation path.""" + path = path.strip() + + supplied = Path(path) + + if supplied.parts and supplied.parts[0] == "docs": + supplied = Path(*supplied.parts[1:]) + + target = (DOCS / supplied).resolve() + + try: + target.relative_to(DOCS.resolve()) + except ValueError as exc: + raise ValueError("Path must be inside the UltraPlot documentation.") from exc + + return target + + +@mcp.tool() +def ping() -> str: + """Check whether the UltraPlot MCP server is working.""" + log.info("ping") + return "pong" + + +@mcp.tool() +def search_docs( + query: str, + limit: int = 8, +) -> list[dict[str, Any]]: + """ + Search UltraPlot documentation and examples. + + Use this for questions about plotting tasks, UltraPlot behavior, + configuration, layouts, projections, legends, colorbars, plotting + commands, formatting, and usage examples. + + Release notes are intentionally excluded. + """ + log.info("search_docs: %r", query) + + query = query.strip() + + if not query: + return [] + + if not DOCS.is_dir(): + return [ + { + "error": "UltraPlot documentation is not installed.", + "docs_path": str(DOCS), + } + ] + + limit = max(1, min(limit, 20)) + terms = _query_terms(query) + + results: list[dict[str, Any]] = [] + + for path in _text_files(): + try: + text = path.read_text( + encoding="utf-8", + errors="replace", + ) + except OSError: + continue + + score = _score_document( + text, + query, + terms, + path, + ) + + if score <= 0: + continue + + results.append( + { + "path": _display_path(path), + "score": score, + "content": _snippet( + text, + query, + terms, + ), + } + ) + + results.sort( + key=lambda result: result["score"], + reverse=True, + ) + + return results[:limit] + + +@mcp.tool() +def search_release_notes( + query: str, +) -> list[dict[str, Any]]: + """ + Search UltraPlot release notes. + + Use this for questions about when features were introduced, + version history, recent changes, fixes, or deprecations. + """ + log.info("search_release_notes: %r", query) + + path = DOCS / "whats_new.rst" + + if not path.is_file(): + return [] + + try: + text = path.read_text( + encoding="utf-8", + errors="replace", + ) + except OSError: + return [] + + query = query.strip() + + if not query: + return [] + + terms = _query_terms(query) + + score = _score_document( + text, + query, + terms, + path, + ) + + if score <= 0: + return [] + + return [ + { + "path": _display_path(path), + "score": score, + "content": _snippet( + text, + query, + terms, + before=900, + after=3200, + ), + } + ] + + +@mcp.tool() +def get_api(symbol: str) -> dict[str, Any]: + """ + Inspect a live UltraPlot Python object. + + Examples: + ultraplot.subplots + ultraplot.axes.Axes.format + ultraplot.axes.PlotAxes.plot + ultraplot.figure.Figure.colorbar + + The ``ultraplot.`` prefix may be omitted. + """ + log.info("get_api: %r", symbol) + + symbol = symbol.strip() + + # Be forgiving when entered manually in MCP Inspector. + if symbol.startswith("symbol:"): + symbol = symbol.removeprefix("symbol:").strip() + + if not symbol.startswith("ultraplot"): + symbol = f"ultraplot.{symbol}" + + obj = pydoc.locate(symbol) + + if obj is None: + return { + "found": False, + "symbol": symbol, + } + + try: + unwrapped = inspect.unwrap(obj) + except ValueError: + unwrapped = obj + + try: + signature = str(inspect.signature(unwrapped)) + except (TypeError, ValueError): + signature = None + + try: + source_file = inspect.getsourcefile(unwrapped) + except TypeError: + source_file = None + + try: + source_line = inspect.getsourcelines(unwrapped)[1] + except (OSError, TypeError, ValueError): + source_line = None + + try: + docstring = inspect.getdoc(obj) + except Exception: + docstring = None + + return { + "found": True, + "symbol": symbol, + "signature": signature, + "docstring": docstring, + "source_file": source_file, + "source_line": source_line, + } + + +@mcp.tool() +def get_source( + symbol: str, +) -> dict[str, Any]: + """ + Return source code for an UltraPlot Python object. + + Use this when the documentation does not fully explain the + implementation or current behavior. + """ + log.info("get_source: %r", symbol) + + symbol = symbol.strip() + + if not symbol.startswith("ultraplot"): + symbol = f"ultraplot.{symbol}" + + obj = pydoc.locate(symbol) + + if obj is None: + return { + "found": False, + "symbol": symbol, + } + + try: + obj = inspect.unwrap(obj) + except ValueError: + pass + + try: + source = inspect.getsource(obj) + except (OSError, TypeError, ValueError): + source = None + + try: + source_file = inspect.getsourcefile(obj) + except TypeError: + source_file = None + + try: + source_line = inspect.getsourcelines(obj)[1] + except (OSError, TypeError, ValueError): + source_line = None + + return { + "found": True, + "symbol": symbol, + "source_file": source_file, + "source_line": source_line, + "source": source, + } + + +@mcp.tool() +def read_doc(path: str) -> str: + """ + Read an UltraPlot documentation or example file. + + Normally use a path returned by ``search_docs``, for example: + docs/subplots.py + docs/projections.py + docs/why.rst + """ + log.info("read_doc: %r", path) + + target = _resolve_doc_path(path) + + if not target.is_file(): + raise FileNotFoundError(path) + + if target.suffix.lower() not in { + ".py", + ".rst", + ".md", + ".txt", + }: + raise ValueError(f"Unsupported documentation file type: {target.suffix}") + + return target.read_text( + encoding="utf-8", + errors="replace", + ) + + +def _mcp_server_command() -> list[str]: + """ + Return a stable command Codex can use to launch this MCP. + + Prefer the installed console script. During development, fall back to + the current Python interpreter and this module. + """ + executable = shutil.which("ultraplot-mcp") + + if executable: + return [str(Path(executable).resolve())] + + # Useful when testing directly from the source checkout. + return [ + sys.executable, + "-m", + "ultraplot.mcp", + ] + + +def install_codex() -> None: + """Register the UltraPlot MCP server with Codex.""" + codex = shutil.which("codex") + + if codex is None: + raise RuntimeError( + "Codex CLI was not found on PATH.\n" + "Install Codex first, then run:\n\n" + " ultraplot-mcp install codex" + ) + + server_command = _mcp_server_command() + + command = [ + codex, + "mcp", + "add", + "ultraplot", + "--", + *server_command, + ] + + log.info( + "Registering UltraPlot MCP with Codex: %s", + " ".join(command), + ) + + try: + subprocess.run( + command, + check=True, + ) + except subprocess.CalledProcessError as exc: + raise RuntimeError( + "Codex could not register the UltraPlot MCP.\n" + "If an UltraPlot MCP entry already exists, remove or update " + "that entry and try again." + ) from exc + + print() + print("UltraPlot MCP registered with Codex.") + print() + print("Codex will launch:") + print() + print(" " + " ".join(server_command)) + print() + print("Restart Codex, then use /mcp to verify the connection.") + + +def _print_help() -> None: + print("""Usage: + ultraplot-mcp + Run the UltraPlot MCP server over stdio. + + ultraplot-mcp install codex + Register the installed UltraPlot MCP with Codex. + +Examples: + ultraplot-mcp + ultraplot-mcp install codex +""") + + +def main() -> None: + args = sys.argv[1:] + + if not args: + mcp.run() + return + + if args in (["-h"], ["--help"], ["help"]): + _print_help() + return + + if args == ["install", "codex"]: + install_codex() + return + + print( + f"Unknown command: {' '.join(args)}\n", + file=sys.stderr, + ) + _print_help() + raise SystemExit(2) + + +if __name__ == "__main__": + main() diff --git a/ultraplot/tests/test_imports.py b/ultraplot/tests/test_imports.py index ff253a8d7..f731496de 100644 --- a/ultraplot/tests/test_imports.py +++ b/ultraplot/tests/test_imports.py @@ -76,6 +76,39 @@ def test_setup_uses_rc_eager_import(): assert out == "ok" +@pytest.mark.parametrize( + "statement", + [ + "from ultraplot import *", + "uplt.setup(eager=True)", + 'uplt.rc["ultraplot.eager_import"] = True; uplt.setup()', + "uplt.__all__", + ], +) +def test_public_imports_do_not_load_optional_mcp(statement): + code = """ +import importlib.abc +import sys + +class BlockMCP(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "mcp" or fullname.startswith("mcp."): + raise ModuleNotFoundError("MCP is not installed", name=fullname) + +sys.meta_path.insert(0, BlockMCP()) +import ultraplot as uplt +""" + code += statement + "\n" + code += """ +assert "mcp" not in uplt.__all__ +assert "ultraplot.mcp" not in sys.modules +assert "mcp" not in sys.modules +assert callable(uplt.subplots) +print("ok") +""" + assert _run(code) == "ok" + + def test_dir_populates_attr_map(monkeypatch): import ultraplot as uplt diff --git a/ultraplot/tests/test_mcp.py b/ultraplot/tests/test_mcp.py new file mode 100644 index 000000000..01f37ddd5 --- /dev/null +++ b/ultraplot/tests/test_mcp.py @@ -0,0 +1,306 @@ +"""Exercise the optional MCP integration without modifying client configuration.""" + +import asyncio +import importlib +import json +import os +from pathlib import Path +import subprocess +import sys +from unittest.mock import Mock + +import pytest + +sdk = pytest.importorskip("mcp.server") +if not hasattr(sdk, "MCPServer"): + pytest.skip("MCP integration requires SDK 2.x", allow_module_level=True) +server = importlib.import_module("ultraplot.mcp") + + +@pytest.fixture +def docs(tmp_path, monkeypatch): + root = tmp_path / "docs" + root.mkdir() + (root / "colorbars.rst").write_text("Shared colorbars\nUse a shared colorbar.") + (root / "example.py").write_text("# A colorbar example") + (root / "unrelated.md").write_text("Geographic projections") + (root / "whats_new.rst").write_text("Added shared colorbars in version 2.6.") + (root / "changelog.rst").write_text("shared colorbars " * 100) + monkeypatch.setattr(server, "DOCS", root) + return root + + +def test_ping(): + assert server.ping() == "pong" + + +def test_search_ranks_relevant_docs_and_excludes_release_notes(docs): + results = server.search_docs(" Shared colorbars ") + assert results[0]["path"] == "docs/colorbars.rst" + assert "Shared colorbars" in results[0]["content"] + assert all(item["score"] > 0 for item in results) + assert not any( + "changelog" in item["path"] or "whats_new" in item["path"] for item in results + ) + assert server.search_docs("missingword") == [] + assert server.search_docs(" ") == [] + + +@pytest.mark.parametrize("limit, expected", [(-1, 1), (2, 2), (100, 20)]) +def test_search_limits_results(docs, limit, expected): + for i in range(25): + (docs / f"example{i}.md").write_text("colorbar") + assert len(server.search_docs("colorbar", limit=limit)) == expected + + +def test_missing_documentation(monkeypatch, tmp_path): + monkeypatch.setattr(server, "DOCS", tmp_path / "missing") + assert "not installed" in server.search_docs("colorbar")[0]["error"] + assert server.search_release_notes("colorbar") == [] + assert list(server._text_files()) == [] + + +def test_unreadable_documents_are_skipped(docs, monkeypatch): + original = Path.read_text + + def read(path, *args, **kwargs): + if path.suffix == ".rst": + raise OSError("unreadable") + return original(path, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", read) + assert server.search_docs("colorbar")[0]["path"] == "docs/example.py" + assert server.search_release_notes("colorbar") == [] + + +def test_release_notes_search(docs): + result = server.search_release_notes(" shared colorbars ") + assert len(result) == 1 + assert result[0]["path"] == "docs/whats_new.rst" + assert "version 2.6" in result[0]["content"] + assert server.search_release_notes("missingword") == [] + assert server.search_release_notes(" ") == [] + + +def test_search_returns_snippet_around_match(docs): + (docs / "long.md").write_text("x " * 3000 + "needle" + " y" * 3000) + result = server.search_docs("needle")[0] + assert "needle" in result["content"] + assert result["content"].startswith("...\n") + assert result["content"].endswith("\n...") + assert len(result["content"]) < 3100 + + +def test_snippet_prefers_cluster_when_phrase_is_absent(): + text = "alpha " + "x " * 1500 + "beta alpha beta" + snippet = server._snippet( + text, "alpha beta gamma", ["alpha", "beta", "gamma"], before=5, after=30 + ) + assert "beta alpha beta" in snippet + assert server._snippet("plain text", "missing", ["missing"]) == "plain text" + + +@pytest.mark.parametrize( + "query, expected", + [("How do I use Colorbars?", ["colorbars"]), ("a", ["a"]), ("", [])], +) +def test_query_normalization(query, expected): + assert server._query_terms(query) == expected + + +@pytest.mark.parametrize("suffix", [".py", ".rst", ".md", ".txt"]) +def test_read_doc_accepts_supported_files(docs, suffix): + path = docs / ("sample" + suffix) + path.write_text("Example π", encoding="utf-8") + assert server.read_doc(path.name) == "Example π" + assert server.read_doc(" docs/" + path.name + " ") == "Example π" + + +def test_read_doc_rejects_invalid_files(docs): + with pytest.raises(FileNotFoundError): + server.read_doc("missing.rst") + (docs / "image.png").write_bytes(b"image") + with pytest.raises(ValueError, match="Unsupported"): + server.read_doc("image.png") + + +def test_read_doc_rejects_paths_outside_docs(docs): + outside = docs.parent / "secret.txt" + outside.write_text("private") + (docs / "link.txt").symlink_to(outside) + for path in ("../secret.txt", "docs/../secret.txt", str(outside), "link.txt"): + with pytest.raises(ValueError, match="inside"): + server.read_doc(path) + + +@pytest.mark.parametrize( + "symbol", + ["axes.Axes.format", "ultraplot.axes.Axes.format", " symbol: axes.Axes.format "], +) +def test_get_api_inspects_live_object(symbol): + result = server.get_api(symbol) + assert result["found"] + assert result["symbol"] == "ultraplot.axes.Axes.format" + assert result["signature"] + assert "title" in result["docstring"] + assert result["source_file"].endswith("base.py") + assert result["source_line"] > 0 + + +@pytest.mark.parametrize("tool", [server.get_api, server.get_source]) +def test_unknown_symbol(tool): + assert tool("_no_such_symbol_") == { + "found": False, + "symbol": "ultraplot._no_such_symbol_", + } + + +def test_get_source_inspects_live_object(): + result = server.get_source("axes.Axes.format") + assert result["found"] + assert "def format(" in result["source"] + assert result["source_file"].endswith("base.py") + assert result["source_line"] > 0 + + +@pytest.mark.parametrize("tool", [server.get_api, server.get_source]) +def test_non_callable_object_has_no_source(tool): + result = tool("__version__") + assert result["found"] + assert result["source_file"] is None + assert result["source_line"] is None + assert result.get("signature") is None + assert result.get("source") is None + + +@pytest.mark.parametrize("tool", [server.get_api, server.get_source]) +def test_inspection_handles_wrapper_cycle(tool, monkeypatch): + def cyclic(): + """A callable with malformed wrapper metadata.""" + + cyclic.__wrapped__ = cyclic + monkeypatch.setattr(server.pydoc, "locate", lambda symbol: cyclic) + result = tool("cyclic") + assert result["found"] + if tool is server.get_api: + assert result["signature"] is None + else: + assert result["source"] is None + + +def test_get_api_handles_unavailable_docstring(monkeypatch): + monkeypatch.setattr( + server.inspect, "getdoc", Mock(side_effect=RuntimeError("unavailable")) + ) + assert server.get_api("axes.Axes.format")["docstring"] is None + + +@pytest.mark.parametrize("args", [["--help"], ["-h"], ["help"]]) +def test_cli_help(args, monkeypatch, capsys): + monkeypatch.setattr(sys, "argv", ["ultraplot-mcp", *args]) + server.main() + assert "Usage:" in capsys.readouterr().out + + +def test_cli_rejects_unknown_command(monkeypatch, capsys): + monkeypatch.setattr(sys, "argv", ["ultraplot-mcp", "unknown"]) + with pytest.raises(SystemExit) as exc: + server.main() + assert exc.value.code == 2 + assert "Unknown command" in capsys.readouterr().err + + +def test_cli_runs_server(monkeypatch): + run = Mock() + monkeypatch.setattr(server.mcp, "run", run) + monkeypatch.setattr(sys, "argv", ["ultraplot-mcp"]) + server.main() + run.assert_called_once_with() + + +def test_cli_registers_codex(monkeypatch, capsys, tmp_path): + codex = str(tmp_path / "codex") + executable = str(tmp_path / "ultraplot-mcp") + monkeypatch.setattr( + server.shutil, "which", lambda name: codex if name == "codex" else executable + ) + run = Mock() + monkeypatch.setattr(server.subprocess, "run", run) + monkeypatch.setattr(sys, "argv", ["ultraplot-mcp", "install", "codex"]) + server.main() + run.assert_called_once_with( + [codex, "mcp", "add", "ultraplot", "--", executable], check=True + ) + assert "registered" in capsys.readouterr().out + + +def test_registration_requires_codex(monkeypatch): + monkeypatch.setattr(server.shutil, "which", lambda name: None) + with pytest.raises(RuntimeError, match="not found"): + server.install_codex() + + +def test_registration_reports_failure(monkeypatch): + monkeypatch.setattr(server.shutil, "which", lambda name: "/bin/" + name) + monkeypatch.setattr( + server.subprocess, + "run", + Mock(side_effect=subprocess.CalledProcessError(1, ["codex"])), + ) + with pytest.raises(RuntimeError, match="could not register"): + server.install_codex() + + +def test_fallback_launch_command_works(monkeypatch): + monkeypatch.setattr(server.shutil, "which", lambda name: None) + result = subprocess.run( + [*server._mcp_server_command(), "--help"], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert "Usage:" in result.stdout + + +def test_stdio_client_can_discover_and_call_tools(docs): + from mcp import ClientSession, StdioServerParameters + from mcp.client.stdio import stdio_client + + async def exercise(): + params = StdioServerParameters( + command=sys.executable, + args=["-m", "ultraplot.mcp"], + env={**os.environ, "ULTRAPLOT_MCP_DOCS": str(docs)}, + ) + async with stdio_client(params) as (reader, writer): + async with ClientSession(reader, writer) as session: + await session.initialize() + listed = await session.list_tools() + assert {tool.name for tool in listed.tools} == { + "ping", + "search_docs", + "search_release_notes", + "read_doc", + "get_api", + "get_source", + } + for name, arguments, expected in [ + ("ping", {}, "pong"), + ( + "search_docs", + {"query": "shared colorbars"}, + "docs/colorbars.rst", + ), + ("read_doc", {"path": "docs/colorbars.rst"}, "Shared colorbars"), + ("search_release_notes", {"query": "colorbars"}, "version 2.6"), + ("get_api", {"symbol": "axes.Axes.format"}, "title"), + ("get_source", {"symbol": "axes.Axes.format"}, "def format("), + ]: + result = await session.call_tool(name, arguments) + assert not result.is_error, result + assert expected in json.dumps(result.model_dump()) + result = await session.call_tool("read_doc", {"path": "../outside.txt"}) + assert result.is_error + + asyncio.run(asyncio.wait_for(exercise(), timeout=45))