From ee4034d1d47f2b0f5adc4a42076d156132707323 Mon Sep 17 00:00:00 2001 From: Raz Date: Sat, 28 Feb 2026 09:31:07 +0200 Subject: [PATCH] feat(mcp): add source coverage status and search diagnostics --- src/unity_docs_mcp/mcp_server.py | 91 ++++++++++++++++++++++++++++---- src/unity_docs_mcp/tools/ops.py | 30 +++++++++++ tests/test_mcp_meta.py | 39 ++++++++++++++ 3 files changed, 149 insertions(+), 11 deletions(-) diff --git a/src/unity_docs_mcp/mcp_server.py b/src/unity_docs_mcp/mcp_server.py index 8b38038..9f2bffa 100644 --- a/src/unity_docs_mcp/mcp_server.py +++ b/src/unity_docs_mcp/mcp_server.py @@ -6,7 +6,7 @@ import os import sys from pathlib import Path -from typing import List, Optional +from typing import Any, List, Optional from mcp.server.fastmcp import FastMCP @@ -76,17 +76,17 @@ def _response_meta(docstore: DocStore, baked_manifest: Optional[dict] = None) -> return meta -@app.tool() -def search( - query: str, - k: int = 6, - source_types: Optional[List[str] | str] = None, -) -> List[dict]: - docstore = _get_docstore() - meta = _response_meta(docstore) +def _parse_source_types(source_types: Optional[List[str] | str]) -> Optional[List[str]]: + if source_types is None: + return None if isinstance(source_types, str): - source_types = [s.strip() for s in source_types.split(",") if s.strip()] - results = docstore.search(query=query, k=k, source_types=source_types) + values = [s.strip().lower() for s in source_types.split(",") if s.strip()] + else: + values = [str(s).strip().lower() for s in source_types if str(s).strip()] + return values or None + + +def _serialize_search_results(results: list[Any], meta: dict) -> List[dict]: return [ { "chunk_id": r.chunk_id, @@ -104,6 +104,65 @@ def search( ] +@app.tool() +def search( + query: str, + k: int = 6, + source_types: Optional[List[str] | str] = None, + debug: bool = False, +) -> List[dict] | dict: + docstore = _get_docstore() + meta = _response_meta(docstore) + parsed_source_types = _parse_source_types(source_types) + available_source_types = docstore.available_source_types() + known_source_types = docstore.known_source_types() + + invalid_source_types = [] + unavailable_source_types = [] + if parsed_source_types: + requested_set = set(parsed_source_types) + known_set = set(known_source_types) + available_set = set(available_source_types) + invalid_source_types = sorted(requested_set - known_set) + unavailable_source_types = sorted((requested_set & known_set) - available_set) + + if invalid_source_types or unavailable_source_types: + message_parts = [] + if invalid_source_types: + message_parts.append(f"Unsupported source_types: {', '.join(invalid_source_types)}") + if unavailable_source_types: + message_parts.append(f"Requested source_types not present in this index: {', '.join(unavailable_source_types)}") + return { + "error": "invalid_source_types", + "message": ". ".join(message_parts), + "requested_source_types": parsed_source_types or [], + "invalid_source_types": invalid_source_types, + "unavailable_source_types": unavailable_source_types, + "known_source_types": known_source_types, + "available_source_types": available_source_types, + "results": [], + "meta": meta, + } + + results = docstore.search(query=query, k=k, source_types=parsed_source_types) + serialized = _serialize_search_results(results, meta) + if not debug: + return serialized + return { + "results": serialized, + "meta": meta, + "debug": { + "query": query, + "k": k, + "requested_source_types": parsed_source_types or [], + "available_source_types": available_source_types, + "known_source_types": known_source_types, + "retrieval_mode": meta["retrieval_mode"], + "result_count": len(serialized), + }, + } + + @app.tool() def open( doc_id: Optional[str] = None, @@ -177,6 +236,13 @@ def status() -> dict: paths = make_paths(config) baked_manifest = _read_manifest(paths.baked_dir / "manifest.json") index_manifest = _read_manifest(paths.index_dir / "manifest.json") + available_source_types = docstore.available_source_types() + source_type_counts = docstore.source_type_counts() + coverage_warnings: list[str] = [] + if "scriptref" not in set(available_source_types): + coverage_warnings.append( + "scriptref source type is not present in the loaded corpus/index; API symbol lookup coverage may be incomplete" + ) return { "meta": _response_meta(docstore, baked_manifest=baked_manifest), "paths": vars(config.paths), @@ -184,6 +250,9 @@ def status() -> dict: "embedder": vars(config.index.embedder), "baked_manifest": baked_manifest, "index_manifest": index_manifest, + "available_source_types": available_source_types, + "source_type_counts": source_type_counts, + "coverage_warnings": coverage_warnings, } diff --git a/src/unity_docs_mcp/tools/ops.py b/src/unity_docs_mcp/tools/ops.py index 32ac1ae..29e6881 100644 --- a/src/unity_docs_mcp/tools/ops.py +++ b/src/unity_docs_mcp/tools/ops.py @@ -10,6 +10,8 @@ from unity_docs_mcp.index.search import HybridSearcher from unity_docs_mcp.paths import make_paths +_DEFAULT_SOURCE_TYPES = ("manual", "scriptref") + @dataclass class DocRecord: @@ -26,8 +28,10 @@ def __init__(self, config: Config): self.config = config self.paths = make_paths(config) self.corpus = self._load_corpus(self.paths.baked_dir / "corpus.jsonl") + self._doc_source_type_counts = self._count_source_types(self.corpus.values()) self.link_index = self._load_links(self.paths.baked_dir / "link_graph.jsonl") self.searcher = HybridSearcher(config, self.paths.index_dir) + self._chunk_source_type_counts = self._count_source_types(getattr(self.searcher, "chunk_meta", {}).values()) def _load_corpus(self, path: Path) -> Dict[str, DocRecord]: records: Dict[str, DocRecord] = {} @@ -56,6 +60,19 @@ def _load_links(self, path: Path) -> Dict[str, List[str]]: links.setdefault(row["from_doc_id"], []).append(row["to_doc_id"]) return links + @staticmethod + def _count_source_types(rows) -> Dict[str, int]: + counts: Dict[str, int] = {} + for row in rows: + if isinstance(row, dict): + source_type = row.get("source_type", "") + else: + source_type = getattr(row, "source_type", "") + if not source_type: + continue + counts[source_type] = counts.get(source_type, 0) + 1 + return counts + def open_doc(self, doc_id: Optional[str] = None, path: Optional[str] = None) -> Optional[DocRecord]: target_id = doc_id if not target_id and path: @@ -82,3 +99,16 @@ def related(self, doc_id: str, limit: int = 10) -> List[DocRecord]: def search(self, query: str, k: int = 6, source_types: Optional[List[str]] = None) -> List: return self.searcher.search(query=query, k=k, source_types=source_types) + + def available_source_types(self) -> List[str]: + all_types = set(self._doc_source_type_counts) | set(self._chunk_source_type_counts) + return sorted(all_types) + + def known_source_types(self) -> List[str]: + return sorted(set(_DEFAULT_SOURCE_TYPES) | set(self.available_source_types())) + + def source_type_counts(self) -> Dict[str, Dict[str, int]]: + return { + "docs": dict(sorted(self._doc_source_type_counts.items())), + "chunks": dict(sorted(self._chunk_source_type_counts.items())), + } diff --git a/tests/test_mcp_meta.py b/tests/test_mcp_meta.py index 9f87a2a..ef378cc 100644 --- a/tests/test_mcp_meta.py +++ b/tests/test_mcp_meta.py @@ -7,6 +7,7 @@ class _FakeDocStore: def __init__(self) -> None: self.config = Config() + self._available_source_types = ["manual"] def search(self, query: str, k: int = 6, source_types=None): return [ @@ -57,6 +58,15 @@ def related(self, doc_id: str, limit: int = 10): ) ] + def available_source_types(self): + return list(self._available_source_types) + + def known_source_types(self): + return ["manual", "scriptref"] + + def source_type_counts(self): + return {"docs": {"manual": 1}, "chunks": {"manual": 1}} + def _install_fake_docstore(monkeypatch): fake = _FakeDocStore() @@ -91,6 +101,32 @@ def test_open_includes_meta_even_when_missing(monkeypatch): assert "doc_id" not in missing +def test_search_invalid_source_types_returns_actionable_error(monkeypatch): + _install_fake_docstore(monkeypatch) + result = mcp_server.search("Rigidbody.AddForce", source_types="scripting") + assert result["error"] == "invalid_source_types" + assert result["invalid_source_types"] == ["scripting"] + assert "manual" in result["available_source_types"] + + +def test_search_unavailable_source_types_returns_actionable_error(monkeypatch): + _install_fake_docstore(monkeypatch) + result = mcp_server.search("Rigidbody.AddForce", source_types="scriptref") + assert result["error"] == "invalid_source_types" + assert result["invalid_source_types"] == [] + assert result["unavailable_source_types"] == ["scriptref"] + + +def test_search_debug_returns_results_and_debug_block(monkeypatch): + _install_fake_docstore(monkeypatch) + result = mcp_server.search("IJobParallelFor batch size", k=3, debug=True) + assert "results" in result + assert "debug" in result + assert result["debug"]["query"] == "IJobParallelFor batch size" + assert result["debug"]["result_count"] == 1 + assert result["results"][0]["doc_id"] == "manual/job-system-parallel-for-jobs" + + def test_list_files_and_related_include_meta(monkeypatch): _install_fake_docstore(monkeypatch) files = mcp_server.list_files("*parallel-for*") @@ -118,3 +154,6 @@ def _fake_read_manifest(path): assert status["meta"]["built_on"] == "2026-02-20" assert "paths" in status assert "index_manifest" in status + assert status["available_source_types"] == ["manual"] + assert status["source_type_counts"]["docs"]["manual"] == 1 + assert status["coverage_warnings"]