diff --git a/README.md b/README.md index 6d1c6489..ef8f122e 100644 --- a/README.md +++ b/README.md @@ -225,7 +225,7 @@ Semble runs as an MCP server so agents can search any codebase directly as a nat | Tool | Description | |------|-------------| -| `search` | Search a codebase with a natural-language or code query. Pass `repo` as a local path or an https:// git URL. | +| `search` | Search a codebase with a natural-language or code query. Pass `repo` as a local path or an https:// git URL and `content` as `code`, `docs`, `config`, or `all` (default: `code`). | | `find_related` | Given a file path and line number, return chunks semantically similar to the code at that location. | For per-agent setup instructions, see the [installation docs](docs/installation.md#mcp-server). diff --git a/docs/installation.md b/docs/installation.md index 128d313b..8e96af7b 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -323,7 +323,7 @@ Add to `~/.zcode/cli/config.json` under the nested `mcp.servers` key (or use Set -By default the MCP server indexes only code files. To also index documentation, config, or everything, append `--content docs`, `--content config`, or `--content all` to the server command. For example, in Claude Code: +The MCP server indexes each requested content selection on first use and caches it separately. Searches default to code; append `--content docs`, `--content config`, or `--content all` to the server command to change that default. The `content` argument on an individual MCP search overrides it. For example, in Claude Code: ```bash claude mcp add semble -s user -- uvx --from "semble[mcp]" semble --content all diff --git a/src/semble/cache.py b/src/semble/cache.py index bdf09b53..894f3e7f 100644 --- a/src/semble/cache.py +++ b/src/semble/cache.py @@ -34,10 +34,11 @@ def cache_key(path: str) -> str: return hashlib.new("sha256", data).hexdigest() -def find_index_from_cache_folder(path: str) -> Path: - """Finds an index from a cache folder and a project path.""" +def find_index_from_cache_folder(path: str, content: Sequence[ContentType] = (ContentType.CODE,)) -> Path: + """Find an exact content index in the cache for a project path.""" cache_dir = resolve_cache_folder() / cache_key(path) - return cache_dir / "index" + scope = "-".join(sorted({content_type.value for content_type in content})) + return cache_dir / ("index" if scope == ContentType.CODE.value else f"index-{scope}") def _windows_cache_dir(name: str) -> Path: @@ -89,16 +90,14 @@ def resolve_cache_folder() -> Path: def clear_cache(path: str) -> None: - """Clears the cache for the given path.""" - index_path = find_index_from_cache_folder(path) - if index_path.exists(): - shutil.rmtree(index_path) + """Clear all exact content indexes for the given path.""" + shutil.rmtree(find_index_from_cache_folder(path).parent, ignore_errors=True) def save_index_to_cache(index: "SembleIndex", path: str) -> None: """Save an index to the cache folder if it was freshly built.""" if not index.loaded_from_disk: - index.save(find_index_from_cache_folder(path)) + index.save(find_index_from_cache_folder(path, index.content)) def _metadata_matches(metadata: dict, model_path: str, content: Sequence[ContentType]) -> bool: @@ -120,7 +119,7 @@ def _metadata_matches(metadata: dict, model_path: str, content: Sequence[Content def get_validated_cache(path: str, model_path: str | None, content: Sequence[ContentType]) -> Path | None: """Validates the cache folder and returns the index path.""" - index_path = find_index_from_cache_folder(path) + index_path = find_index_from_cache_folder(path, content) if not index_path.exists(): return None @@ -169,7 +168,7 @@ def load_previous_for_incremental( :return: Previous index state, or None if the cache is unavailable or invalid. """ try: - index_path = find_index_from_cache_folder(path) + index_path = find_index_from_cache_folder(path, content) persistence_path = PersistencePath.from_path(index_path) if persistence_path.non_existing(): return None diff --git a/src/semble/cli.py b/src/semble/cli.py index eddd3df2..466b42ef 100644 --- a/src/semble/cli.py +++ b/src/semble/cli.py @@ -12,7 +12,7 @@ from model2vec.utils import get_package_extras -from semble.cache import cache_key, find_index_from_cache_folder, resolve_cache_folder +from semble.cache import cache_key, resolve_cache_folder, save_index_to_cache from semble.index import SembleIndex from semble.index.types import PersistencePath from semble.installer.agents import AGENTS, IntegrationType @@ -40,12 +40,10 @@ def _build_index(path: str, content: list[ContentType]) -> SembleIndex: def _maybe_save_index(index: SembleIndex, path: str) -> None: """Save the index to the cache folder if it was not loaded from disk.""" - if not index.loaded_from_disk: - try: - cache_folder = find_index_from_cache_folder(path) - index.save(cache_folder) - except Exception as e: - print(f"Error saving index: {e}", file=sys.stderr) + try: + save_index_to_cache(index, path) + except Exception as e: + print(f"Error saving index: {e}", file=sys.stderr) def _add_content_args(p: argparse.ArgumentParser) -> None: @@ -56,7 +54,7 @@ def _add_content_args(p: argparse.ArgumentParser) -> None: default=["code"], choices=[ct.value for ct in ContentType] + ["all"], metavar="TYPE", - help="Content types to index (space-separated, e.g. --content code docs). Choices: code, docs, config, all. Default: code.", + help="Content types to search (space-separated, e.g. --content code docs). Choices: code, docs, config, all. Default: code.", ) p.add_argument( "--include-text-files", @@ -146,19 +144,18 @@ def _run_find_related( def _clear_indexes(cache_folder: Path) -> None: """Remove all valid index entries from the cache folder.""" - indexes = [] - for path in cache_folder.glob("*/index"): + indexes: set[Path] = set() + for path in cache_folder.glob("*/index*"): if not _SHA_256_REGEX.match(path.parent.name): continue if PersistencePath.from_path(path).non_existing(): continue - indexes.append(path) + indexes.add(path.parent) if not indexes: print(f"No indexes found to clear in `{cache_folder}`") else: - for path in indexes: - index_folder = path.parent + for index_folder in indexes: rmtree(index_folder) print(f"Cleared index at `{index_folder}`") @@ -175,8 +172,8 @@ def _clear_savings(cache_folder: Path) -> None: def _clear_orphans(cache_folder: Path) -> None: """Remove index entries whose local root_path no longer exists.""" - orphans = [] - for path in cache_folder.glob("*/index"): + orphans: dict[Path, str] = {} + for path in cache_folder.glob("*/index*"): if not _SHA_256_REGEX.match(path.parent.name): continue try: @@ -189,12 +186,12 @@ def _clear_orphans(cache_folder: Path) -> None: if not isinstance(root_path, str) or not root_path or cache_key(root_path) != path.parent.name: continue if not Path(root_path).exists(): - orphans.append((path.parent, root_path)) + orphans[path.parent] = root_path if not orphans: print("No orphaned indexes found") else: - for index_folder, root_path in orphans: + for index_folder, root_path in orphans.items(): rmtree(index_folder) print(f"Cleared orphaned index for `{root_path}`") diff --git a/src/semble/index/index.py b/src/semble/index/index.py index 9899a329..cc199edd 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -124,6 +124,11 @@ def stats(self) -> IndexStats: languages=dict(language_counts), ) + @property + def content(self) -> tuple[ContentType, ...]: + """Return the content types covered by this index.""" + return self._content + @classmethod def from_path( cls, diff --git a/src/semble/installer/agents.py b/src/semble/installer/agents.py index 10f475c2..eaca05ea 100644 --- a/src/semble/installer/agents.py +++ b/src/semble/installer/agents.py @@ -92,7 +92,7 @@ def semble_pin() -> str: Use `mcp__semble__search` to find where something is implemented — instead of using Grep or Glob to discover files. After semble returns the file and line, navigate there directly and read that file. Do not grep for the same content again. -Pass `--content docs` to search documentation and prose, `--content config` for config files, or `--content all` to search code, docs, and config together. +Pass `content="docs"` to the MCP search tool for documentation and prose, `content="config"` for config files, or `content="all"` for everything. On the CLI, use `--content docs`, `--content config`, or `--content all` instead. For CLI fallback or sub-agents without MCP access, use: @@ -111,8 +111,8 @@ def semble_pin() -> str: 1. Call `mcp__semble__search` with a query describing what the code does or its name. The tool returns results with 10 lines of context each (function/class signature + first body lines, enough to confirm the location). 2. Navigate directly to the top result's file and line. Read only the function or class at that location. 3. Make the edit. Do not re-search or grep for the same content. -4. Use `--content docs` for documentation, `--content config` for config files, or `--content all` for everything. -5. Optionally use `mcp__semble__find_related` with `file_path` and `line` to discover similar code elsewhere. +4. Set the MCP search tool's `content` field to `docs`, `config`, or `all` when searching beyond code. +5. Optionally use `mcp__semble__find_related` with `file_path`, `line`, and the same `content` selection to discover similar code elsewhere. 6. Use Grep only when you need every occurrence of a literal string across the whole repo (e.g., all callers of a renamed function). {SEMBLE_END} """ diff --git a/src/semble/mcp.py b/src/semble/mcp.py index 1274b352..b27055ef 100644 --- a/src/semble/mcp.py +++ b/src/semble/mcp.py @@ -7,7 +7,7 @@ from collections import OrderedDict from collections.abc import Sequence from pathlib import Path -from typing import Annotated +from typing import Annotated, Literal from mcp.server.fastmcp import FastMCP from pydantic import Field @@ -27,22 +27,32 @@ _CACHE_MAX_SIZE = 10 # Max number of cached indexes to keep in memory _MIN_REVALIDATE_FACTOR = 3 # Don't recheck staleness sooner than this many times the last build's duration +ContentSelection = Literal["code", "docs", "config", "all"] +_CacheKey = tuple[str, tuple[ContentType, ...]] -async def _get_index( - repo: str, - cache: _IndexCache, -) -> SembleIndex: +async def _get_index(repo: str, cache: _IndexCache, content: Sequence[ContentType]) -> SembleIndex: """Return a cached index for a repo, rejecting unsafe git transport schemes.""" if is_git_url(repo) and not repo.startswith(("https://", "http://")): raise ValueError(f"Only https://, http://, or local directory paths are accepted as `repo`. Got: {repo!r}") try: - return await cache.get(repo) + return await cache.get(repo, content=content) except Exception as exc: raise ValueError(f"Failed to index {repo!r}: {exc}") from exc -def create_server(cache: _IndexCache) -> FastMCP: +def _resolve_content_selection( + content: ContentSelection | None, default_content: Sequence[ContentType] +) -> tuple[ContentType, ...]: + """Resolve an MCP content selection to exact index content types.""" + if content is None: + return tuple(default_content) + if content == "all": + return tuple(ContentType) + return (ContentType(content),) + + +def create_server(cache: _IndexCache, default_content: Sequence[ContentType] = (ContentType.CODE,)) -> FastMCP: """Build and return a configured FastMCP server backed by the given cache.""" server = FastMCP( "semble", @@ -74,6 +84,10 @@ async def search( ge=0, ), ] = 10, + content: Annotated[ + ContentSelection | None, + Field(description="Content to search. Defaults to the MCP server's configured content."), + ] = None, ) -> str: """Search once with a focused query describing what the code does or its name. @@ -81,8 +95,9 @@ async def search( Returns file paths and line numbers — navigate directly there, do not repeat the search. Pass a git URL or local path as `repo`; indexes are cached for the session. """ + selected_content = _resolve_content_selection(content, default_content) try: - index = await _get_index(repo, cache) + index = await _get_index(repo, cache, selected_content) except ValueError as exc: return str(exc) results = index.search(query, top_k=top_k, max_snippet_lines=max_snippet_lines) @@ -109,6 +124,10 @@ async def find_related( ge=0, ), ] = 10, + content: Annotated[ + ContentSelection | None, + Field(description="Content containing the related file. Defaults to the MCP server configuration."), + ] = None, ) -> str: """Find code similar to a known location. @@ -116,8 +135,9 @@ async def find_related( or all tests for a class. Use after `search` when you need related code beyond the primary result. Pass `file_path` and `line` from a prior search result. """ + selected_content = _resolve_content_selection(content, default_content) try: - index = await _get_index(repo, cache) + index = await _get_index(repo, cache, selected_content) except ValueError as exc: return str(exc) chunk = resolve_chunk(index.chunks, file_path, line) @@ -139,7 +159,7 @@ async def serve( content: Sequence[ContentType] = (ContentType.CODE,), ) -> None: """Start an MCP stdio server.""" - cache = _IndexCache(content=content) + cache = _IndexCache() async def _load_and_prewarm() -> None: """Pre-load the embedding model in parallel with starting the server.""" @@ -153,7 +173,7 @@ async def _load_and_prewarm() -> None: cache._model_ready.set() init_task = asyncio.create_task(_load_and_prewarm()) - server = create_server(cache) + server = create_server(cache, default_content=content) try: await server.run_stdio_async() finally: @@ -164,14 +184,13 @@ async def _load_and_prewarm() -> None: class _IndexCache: """Cache of indexed repos and local paths for the lifetime of the MCP server process.""" - def __init__(self, content: Sequence[ContentType] = (ContentType.CODE,)) -> None: + def __init__(self) -> None: """Initialise an empty cache.""" self._model_path: str | None = None self._model_error: BaseException | None = None self._model_ready = asyncio.Event() - self._content = content - self._tasks: OrderedDict[str, asyncio.Task[SembleIndex]] = OrderedDict() # ordered for LRU eviction - self._revalidate_after: dict[str, float] = {} # cache_key -> monotonic time, staleness check is gated until + self._tasks: OrderedDict[_CacheKey, asyncio.Task[SembleIndex]] = OrderedDict() # ordered for LRU eviction + self._revalidate_after: dict[_CacheKey, float] = {} async def _await_model(self) -> str: """Block until the model is installed; re-raise the load error if it failed.""" @@ -181,43 +200,51 @@ async def _await_model(self) -> str: assert self._model_path is not None return self._model_path - def _compute_cache_key(self, source: str, ref: str | None = None) -> str: - """Compute the canonical cache key for a source.""" + def _compute_cache_key( + self, + source: str, + ref: str | None = None, + content: Sequence[ContentType] = (ContentType.CODE,), + ) -> _CacheKey: + """Compute the canonical key for an exact index variant.""" is_git = is_git_url(source) - return (f"{source}@{ref}" if ref else source) if is_git else str(Path(source).resolve()) + source_key = (f"{source}@{ref}" if ref else source) if is_git else str(Path(source).resolve()) + normalized = tuple(content_type for content_type in ContentType if content_type in content) + return source_key, normalized - def _build_and_cache_index(self, source: str, ref: str | None, model_path: str, cache_key: str) -> SembleIndex: + def _build_index(self, source: str, ref: str | None, model_path: str, cache_key: _CacheKey) -> SembleIndex: """Build an index for the given source and cache it.""" + source_key, content = cache_key index = ( - SembleIndex.from_git(source, ref=ref, model_path=model_path, content=self._content) + SembleIndex.from_git(source, ref=ref, model_path=model_path, content=content) if is_git_url(source) - else SembleIndex.from_path(cache_key, model_path=model_path, content=self._content) + else SembleIndex.from_path(source_key, model_path=model_path, content=content) ) try: - save_index_to_cache(index, cache_key) + save_index_to_cache(index, source_key) except Exception: - logger.warning("Failed to save index cache for %r", cache_key, exc_info=True) + logger.warning("Failed to save index cache for %r", source_key, exc_info=True) return index - async def _build_and_track(self, source: str, ref: str | None, model_path: str, cache_key: str) -> SembleIndex: + async def _build_tracked(self, source: str, ref: str | None, model_path: str, cache_key: _CacheKey) -> SembleIndex: """Build an index and, for local paths, record when its staleness cooldown ends. The cooldown write happens after the await, i.e. back on the event loop thread, - regardless of which thread `_build_and_cache_index` itself ran on. + regardless of which thread `_build_index` itself ran on. """ start = time.monotonic() - index = await asyncio.to_thread(self._build_and_cache_index, source, ref, model_path, cache_key) + index = await asyncio.to_thread(self._build_index, source, ref, model_path, cache_key) if not is_git_url(source): finished = time.monotonic() self._revalidate_after[cache_key] = finished + (finished - start) * _MIN_REVALIDATE_FACTOR return index - def evict(self, source: str) -> None: - cache_key = self._compute_cache_key(source) + def evict(self, cache_key: _CacheKey) -> None: + """Evict one exact index variant from memory.""" self._tasks.pop(cache_key, None) self._revalidate_after.pop(cache_key, None) - async def _evict_if_stale(self, source: str, cache_key: str) -> None: + async def _evict_if_stale(self, cache_key: _CacheKey) -> None: """Evict a cached local-path entry whose on-disk cache no longer matches its files. Skipped while inside the cooldown window so repos that are slow to build aren't @@ -226,7 +253,7 @@ async def _evict_if_stale(self, source: str, cache_key: str) -> None: cached = self._tasks.get(cache_key) if ( cached is None - or is_git_url(source) + or is_git_url(cache_key[0]) or not cached.done() or cached.cancelled() or cached.exception() is not None @@ -234,19 +261,24 @@ async def _evict_if_stale(self, source: str, cache_key: str) -> None: return if time.monotonic() < self._revalidate_after.get(cache_key, 0.0): return - validated = await asyncio.to_thread(get_validated_cache, cache_key, self._model_path, self._content) + validated = await asyncio.to_thread(get_validated_cache, cache_key[0], self._model_path, cache_key[1]) # Only evict if this entry hasn't already been replaced by a concurrent caller. if validated is None and self._tasks.get(cache_key) is cached: - self.evict(source) - - async def get(self, source: str, ref: str | None = None) -> SembleIndex: + self.evict(cache_key) + + async def get( + self, + source: str, + ref: str | None = None, + content: Sequence[ContentType] = (ContentType.CODE,), + ) -> SembleIndex: """Return an index for the requested source, building and caching it on first access. Local paths are revalidated against the on-disk cache on every call (subject to a cooldown scaled by build time), so an entry is rebuilt once its files change. """ - cache_key = self._compute_cache_key(source, ref) - await self._evict_if_stale(source, cache_key) + cache_key = self._compute_cache_key(source, ref, content) + await self._evict_if_stale(cache_key) if cache_key not in self._tasks: model_path = await self._await_model() @@ -255,17 +287,17 @@ async def get(self, source: str, ref: str | None = None) -> SembleIndex: if len(self._tasks) >= _CACHE_MAX_SIZE: evicted_key, _ = self._tasks.popitem(last=False) self._revalidate_after.pop(evicted_key, None) - self._tasks[cache_key] = asyncio.create_task(self._build_and_track(source, ref, model_path, cache_key)) + self._tasks[cache_key] = asyncio.create_task(self._build_tracked(source, ref, model_path, cache_key)) self._tasks.move_to_end(cache_key) task = self._tasks[cache_key] try: return await asyncio.shield(task) except asyncio.CancelledError: # pragma: no cover if task.done(): - self.evict(source) + self.evict(cache_key) raise except Exception: # Only evict if this task hasn't already been replaced by evict()+get(). if self._tasks.get(cache_key) is task: - self.evict(source) + self.evict(cache_key) raise diff --git a/src/semble/version.py b/src/semble/version.py index 5073cce7..1ae5ed05 100644 --- a/src/semble/version.py +++ b/src/semble/version.py @@ -1,2 +1,2 @@ -__version_triple__ = (0, 5, 4) +__version_triple__ = (0, 5, 5) __version__ = ".".join(map(str, __version_triple__)) diff --git a/tests/index/test_index.py b/tests/index/test_index.py index e80319bf..0f20c663 100644 --- a/tests/index/test_index.py +++ b/tests/index/test_index.py @@ -86,6 +86,7 @@ def test_tiny_invalid_utf8_file_status_does_not_crash(tmp_path: Path) -> None: def test_index_language_counts(indexed_index: SembleIndex) -> None: """Language breakdown in stats includes python with at least one chunk.""" stats = indexed_index.stats + assert indexed_index.content == (ContentType.CODE,) assert "python" in stats.languages assert stats.languages["python"] > 0 diff --git a/tests/test_cache.py b/tests/test_cache.py index 294ebfa5..e9c197d3 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -24,10 +24,12 @@ def test_find_index_from_cache_folder_local_path(tmp_path: Path) -> None: - """Local paths are normalised before hashing, result ends with /index.""" + """Local paths are normalised and content variants get distinct index directories.""" result = find_index_from_cache_folder(str(tmp_path)) assert result.name == "index" assert result == find_index_from_cache_folder(str(tmp_path)) + assert find_index_from_cache_folder(str(tmp_path), [ContentType.DOCS]).name == "index-docs" + assert find_index_from_cache_folder(str(tmp_path), list(ContentType)).name == "index-code-config-docs" def test_find_index_from_cache_folder_git_url() -> None: @@ -74,9 +76,10 @@ def test_cache_dir_no_env(fn: object, expected_rel: Path) -> None: def test_save_index_to_cache(tmp_path: Path) -> None: """A freshly built index is saved under its cache key.""" - index = MagicMock(loaded_from_disk=False) - with patch("semble.cache.find_index_from_cache_folder", return_value=tmp_path / "index"): + index = MagicMock(loaded_from_disk=False, content=(ContentType.DOCS,)) + with patch("semble.cache.find_index_from_cache_folder", return_value=tmp_path / "index") as mock_find: save_index_to_cache(index, "repo") + mock_find.assert_called_once_with("repo", (ContentType.DOCS,)) index.save.assert_called_once_with(tmp_path / "index") @@ -119,14 +122,18 @@ def test_resolve_cache_folder_semble_cache_location(tmp_path: Path) -> None: def test_clear_cache(tmp_path: Path) -> None: - """clear_cache removes the index directory when it exists and is a no-op otherwise.""" - index_path = tmp_path / "index" + """clear_cache removes every content variant and is a no-op when none exist.""" + cache_dir = tmp_path / "repo" + index_path = cache_dir / "index" with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): clear_cache("/some/path") # no-op: path doesn't exist yet - index_path.mkdir() + index_path.mkdir(parents=True) + docs_path = cache_dir / "index-docs" + docs_path.mkdir() with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): clear_cache("/some/path") assert not index_path.exists() + assert not docs_path.exists() def _write_metadata( diff --git a/tests/test_cli.py b/tests/test_cli.py index e4245f60..029c4079 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -224,11 +224,9 @@ def test_cli_content_argument( def test_maybe_save_index_logs_error_on_save_failure(capsys: pytest.CaptureFixture[str]) -> None: - """_maybe_save_index prints to stderr when index.save raises.""" + """_maybe_save_index prints to stderr when cache persistence fails.""" fake_index = MagicMock() - fake_index.loaded_from_disk = False - fake_index.save.side_effect = OSError("disk full") - with patch("semble.cli.find_index_from_cache_folder", return_value=Path("/cache")): + with patch("semble.cli.save_index_to_cache", side_effect=OSError("disk full")): _maybe_save_index(fake_index, "/some/path") assert "Error saving index" in capsys.readouterr().err @@ -242,9 +240,11 @@ def test_agent_file_tools_are_bash_only() -> None: assert not any("mcp__" in t for t in tools) -def _make_valid_index_dir(cache_folder: Path, sha: str = "a" * 64, metadata: str = "{}") -> Path: +def _make_valid_index_dir( + cache_folder: Path, sha: str = "a" * 64, metadata: str = "{}", index_name: str = "index" +) -> Path: """Create a fake valid index directory with the expected structure.""" - index_dir = cache_folder / sha / "index" + index_dir = cache_folder / sha / index_name index_dir.mkdir(parents=True) # Create the files that PersistencePath.non_existing checks (index_dir / "chunks.json").write_text("[]") @@ -269,7 +269,7 @@ def test_run_clear_index( """_run_clear('index') finds valid indexes, and skips non-SHA/incomplete/empty dirs.""" if scenario == "valid": _make_valid_index_dir(tmp_path, "a" * 64) - _make_valid_index_dir(tmp_path, "b" * 64) + _make_valid_index_dir(tmp_path, "b" * 64, index_name="index-docs") elif scenario == "non_sha": bad_dir = tmp_path / "not-a-sha" / "index" bad_dir.mkdir(parents=True) @@ -305,7 +305,7 @@ def test_run_clear_orphans(scenario: str, tmp_path: Path, capsys: pytest.Capture root.mkdir() sha = hashlib.sha256(str(root.resolve()).encode("utf-8")).hexdigest() if scenario == "orphan": - _make_valid_index_dir(cache_folder, sha, metadata=json.dumps({"root_path": str(root)})) + _make_valid_index_dir(cache_folder, sha, metadata=json.dumps({"root_path": str(root)}), index_name="index-docs") root.rmdir() elif scenario == "live": _make_valid_index_dir(cache_folder, sha, metadata=json.dumps({"root_path": str(root)})) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 0f318f26..6ef5c0dc 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import asyncio +import json import threading import time from pathlib import Path @@ -9,7 +10,7 @@ from model2vec import StaticModel from semble.mcp import _CACHE_MAX_SIZE, _IndexCache, create_server, serve -from semble.types import Chunk, SearchResult +from semble.types import Chunk, ContentType, SearchResult from semble.utils import format_results, is_git_url, resolve_chunk from tests.conftest import make_chunk @@ -144,10 +145,17 @@ async def test_index_cache_builds_and_caches( ): first = await cache.get(resolved_source) second = await cache.get(resolved_source) + docs_first = await cache.get(resolved_source, content=(ContentType.DOCS,)) + docs_second = await cache.get(resolved_source, content=(ContentType.DOCS,)) assert first is fake_index assert second is fake_index - mock_build.assert_called_once() - mock_save.assert_called_once_with(fake_index, cache._compute_cache_key(resolved_source)) + assert docs_first is fake_index + assert docs_second is fake_index + assert [call.kwargs["content"] for call in mock_build.call_args_list] == [ + (ContentType.CODE,), + (ContentType.DOCS,), + ] + assert mock_save.call_count == 2 @pytest.mark.anyio @@ -186,12 +194,12 @@ async def test_index_cache_staleness_check_scope( @pytest.mark.anyio async def test_index_cache_skips_staleness_check_during_cooldown(cache: _IndexCache, tmp_path: Path) -> None: """A slow-to-build local path is not revalidated again until its cooldown elapses.""" - cache_key = str(tmp_path.resolve()) + cache_key = cache._compute_cache_key(str(tmp_path)) cache._tasks[cache_key] = asyncio.create_task(_succeed()) await asyncio.sleep(0) # let the task finish cache._revalidate_after[cache_key] = time.monotonic() + 30.0 # a build that took 10s, just finished with patch("semble.mcp.get_validated_cache") as mock_validate: - await cache._evict_if_stale(str(tmp_path), cache_key) + await cache._evict_if_stale(cache_key) mock_validate.assert_not_called() @@ -206,17 +214,18 @@ async def test_index_cache_skips_staleness_check_for_failed_task(cache: _IndexCa async def _raise() -> MagicMock: raise RuntimeError("boom") - cache._tasks[str(tmp_path.resolve())] = asyncio.create_task(_raise()) + cache_key = cache._compute_cache_key(str(tmp_path)) + cache._tasks[cache_key] = asyncio.create_task(_raise()) await asyncio.sleep(0) # let the task finish with patch("semble.mcp.get_validated_cache") as mock_validate: - await cache._evict_if_stale(str(tmp_path), str(tmp_path.resolve())) + await cache._evict_if_stale(cache_key) mock_validate.assert_not_called() @pytest.mark.anyio async def test_index_cache_does_not_evict_entry_replaced_during_validation(cache: _IndexCache, tmp_path: Path) -> None: """If a concurrent caller already replaced a stale entry, _evict_if_stale must not evict the new one.""" - cache_key = str(tmp_path.resolve()) + cache_key = cache._compute_cache_key(str(tmp_path)) cache._tasks[cache_key] = asyncio.create_task(_succeed()) await asyncio.sleep(0) cache._revalidate_after[cache_key] = 0.0 # cooldown already elapsed @@ -229,7 +238,7 @@ def _replace_entry_then_report_stale(*args: object, **kwargs: object) -> None: return None with patch("semble.mcp.get_validated_cache", side_effect=_replace_entry_then_report_stale): - await cache._evict_if_stale(str(tmp_path), cache_key) + await cache._evict_if_stale(cache_key) assert cache._tasks.get(cache_key) is replacement_task @@ -348,6 +357,35 @@ async def test_tool_output( assert substring in text +@pytest.mark.anyio +async def test_search_builds_exact_content_indexes( + cache: _IndexCache, + mock_model: StaticModel, + tmp_project: Path, +) -> None: + """MCP search lazily builds the exact requested content index.""" + (tmp_project / "settings.toml").write_text("project = 'semble'\n") + expected = [ + (None, {".py"}), + ("docs", {".md"}), + ("config", {".toml"}), + ("all", {".md", ".py", ".toml"}), + ] + + with ( + patch("semble.index.index.load_model", return_value=(mock_model, "/fake/model")), + patch("semble.mcp.save_index_to_cache"), + ): + server = create_server(cache) + for content, expected_suffixes in expected: + args = {"query": "project", "repo": str(tmp_project), "top_k": 20} + if content is not None: + args["content"] = content + result = await server.call_tool("search", args) + payload = json.loads(_tool_text(result)) + assert {Path(item["file_path"]).suffix for item in payload["results"]} == expected_suffixes + + @pytest.mark.anyio @pytest.mark.parametrize( ("load_err", "stdio_yields"), @@ -458,7 +496,7 @@ async def test_index_cache_lru_eviction(cache: _IndexCache, tmp_path: Path) -> N with patch("semble.mcp.SembleIndex.from_path", return_value=MagicMock()): for d in dirs[:_CACHE_MAX_SIZE]: await cache.get(str(d)) - first_key = str(dirs[0].resolve()) + first_key = cache._compute_cache_key(str(dirs[0])) assert first_key in cache._tasks await cache.get(str(dirs[_CACHE_MAX_SIZE])) assert first_key not in cache._tasks @@ -466,13 +504,13 @@ async def test_index_cache_lru_eviction(cache: _IndexCache, tmp_path: Path) -> N def test_cache_evict(cache: _IndexCache, tmp_path: Path) -> None: - """evict() removes an existing cache entry by resolved path.""" - key = str(tmp_path.resolve()) + """evict() removes an existing exact cache entry.""" + key = cache._compute_cache_key(str(tmp_path)) cache._tasks[key] = MagicMock() - cache.evict(str(tmp_path)) + cache.evict(key) assert key not in cache._tasks def test_cache_evict_missing(cache: _IndexCache, tmp_path: Path) -> None: - """evict() on an unknown path is a no-op.""" - cache.evict(str(tmp_path)) # should not raise + """evict() on an unknown key is a no-op.""" + cache.evict(cache._compute_cache_key(str(tmp_path))) # should not raise