diff --git a/docs/installation.md b/docs/installation.md index fd4d74ed..07f05535 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -329,6 +329,8 @@ The MCP server indexes each requested content selection on first use and caches claude mcp add semble -s user -- uvx --from "semble[mcp]" semble --content all ``` +Indexes stay in memory for the lifetime of the server. To free memory while idle, set `SEMBLE_MCP_IDLE_TIMEOUT` to a number of seconds; indexes unused for that long are dropped from memory and reloaded from the disk cache on the next search. + ### Instructions (AGENTS.md / CLAUDE.md) Add the snippet below to your `AGENTS.md` or `CLAUDE.md` so your agent knows when and how to call the semble CLI: diff --git a/src/semble/mcp.py b/src/semble/mcp.py index f03f7435..13e35704 100644 --- a/src/semble/mcp.py +++ b/src/semble/mcp.py @@ -3,6 +3,7 @@ import asyncio import json import logging +import os import time from collections import OrderedDict from collections.abc import Sequence @@ -27,6 +28,9 @@ ) _CACHE_MAX_SIZE = 10 # Max number of cached indexes to keep in memory +_CACHE_IDLE_TIMEOUT = float( + os.environ.get("SEMBLE_MCP_IDLE_TIMEOUT", 0) +) # Idle seconds before dropping an index (0 = never) _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, ...]] @@ -198,6 +202,7 @@ class _IndexCache: def __init__(self) -> None: """Initialise an empty cache.""" + self._idle_timers: dict[_CacheKey, asyncio.TimerHandle] = {} self._model_path: str | None = None self._model_error: BaseException | None = None self._model_ready = asyncio.Event() @@ -263,6 +268,13 @@ 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) + if (timer := self._idle_timers.pop(cache_key, None)) is not None: + timer.cancel() + + def _evict_idle(self, cache_key: _CacheKey) -> None: + """Drop an entry that has not been accessed within the idle timeout; the disk cache is kept.""" + self.evict(cache_key) + self._merged = None # may hold a reference to the evicted index 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. @@ -305,13 +317,12 @@ async def get( # Re-check after the await: another caller may have populated the entry. if cache_key not in self._tasks: if len(self._tasks) >= _CACHE_MAX_SIZE: - evicted_key, _ = self._tasks.popitem(last=False) - self._revalidate_after.pop(evicted_key, None) + self.evict(next(iter(self._tasks))) 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) + index = await asyncio.shield(task) except asyncio.CancelledError: # pragma: no cover if task.done(): self.evict(cache_key) @@ -321,3 +332,11 @@ async def get( if self._tasks.get(cache_key) is task: self.evict(cache_key) raise + # Start the idle timer only once the index is ready, so slow builds are not evicted mid-build. + if _CACHE_IDLE_TIMEOUT > 0 and self._tasks.get(cache_key) is task: + if timer := self._idle_timers.get(cache_key): + timer.cancel() + self._idle_timers[cache_key] = asyncio.get_running_loop().call_later( + _CACHE_IDLE_TIMEOUT, self._evict_idle, cache_key + ) + return index diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 22acf230..faa81f7d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -523,6 +523,33 @@ async def test_index_cache_lru_eviction(cache: _IndexCache, tmp_path: Path) -> N assert len(cache._tasks) == _CACHE_MAX_SIZE +@pytest.mark.anyio +async def test_index_cache_idle_timeout_eviction(cache: _IndexCache, tmp_path: Path) -> None: + """Entries are dropped once unused for the idle timeout, which starts after the build and resets on access.""" + + def slow_build(*args: Any, **kwargs: Any) -> MagicMock: + time.sleep(0.15) # longer than the timeout + return MagicMock() + + key = cache._compute_cache_key(str(tmp_path)) + with ( + patch("semble.mcp._CACHE_IDLE_TIMEOUT", 0.1), + patch("semble.mcp.SembleIndex.from_path", side_effect=slow_build), + patch("semble.mcp.get_validated_cache", return_value=MagicMock()), + ): + await cache.get(str(tmp_path)) + assert key in cache._tasks + first_timer = cache._idle_timers[key] + cache._merged = ([], MagicMock()) + await cache.get(str(tmp_path)) + assert first_timer.cancelled() + assert key in cache._tasks + await asyncio.sleep(0.3) + assert key not in cache._tasks + assert key not in cache._idle_timers + assert cache._merged is None + + def test_cache_evict(cache: _IndexCache, tmp_path: Path) -> None: """evict() removes an existing exact cache entry.""" key = cache._compute_cache_key(str(tmp_path))