coderay 1.0.0__py3-none-any.whl

coderay/mcp_server/server.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+logger = logging.getLogger(__name__)
+
+mcp = FastMCP("coderay")
+
+DEFAULT_INDEX_DIR = ".index"
+
+_retrieval_cache: dict[Path, Any] = {}
+_state_machine_cache: dict[Path, Any] = {}
+
+
+def _resolve_index_dir(index_dir: str | None = None) -> Path:
+    """Resolve the index directory to an absolute path."""
+    return Path(index_dir or DEFAULT_INDEX_DIR).resolve()
+
+
+def _get_retrieval(index_dir: str | None = None):
+    """Return a cached Retrieval instance for the given index directory."""
+    idx = _resolve_index_dir(index_dir)
+    if idx not in _retrieval_cache:
+        from coderay.retrieval.search import Retrieval
+
+        _retrieval_cache[idx] = Retrieval(idx)
+    return _retrieval_cache[idx]
+
+
+def _load_graph(index_dir: str | None = None):
+    """Load the code graph from disk, or return None if absent."""
+    from coderay.graph.builder import load_graph
+
+    return load_graph(_resolve_index_dir(index_dir))
+
+
+def _get_state_machine(index_dir: str | None = None):
+    """Return a cached StateMachine instance for the given index directory."""
+    idx = _resolve_index_dir(index_dir)
+    if idx not in _state_machine_cache:
+        from coderay.state.machine import StateMachine
+
+        _state_machine_cache[idx] = StateMachine(idx)
+    return _state_machine_cache[idx]
+
+
+def _load_state(index_dir: str | None = None):
+    """Load the current IndexMeta state, or None if no run has completed."""
+    return _get_state_machine(index_dir).current_state
+
+
+@mcp.tool()
+def semantic_search(
+    query: str,
+    top_k: int = 10,
+    path_prefix: str | None = None,
+    language: str | None = None,
+    index_dir: str | None = None,
+) -> str:
+    """Search code by meaning."""
+    retrieval = _get_retrieval(index_dir)
+    state = _load_state(index_dir)
+    if state is None:
+        return json.dumps({"error": "No index state found. Run 'coderay build' first."})
+    try:
+        results = retrieval.search(
+            query,
+            state,
+            top_k=top_k,
+            path_prefix=path_prefix,
+            language=language,
+        )
+    except RuntimeError as e:
+        return json.dumps({"error": str(e)})
+    score_type = results[0].get("score_type", "cosine") if results else "cosine"
+    return json.dumps(
+        {
+            "score_type": score_type,
+            "score_description": (
+                "cosine similarity (0-1, higher = more similar)"
+                if score_type == "cosine"
+                else "RRF rank fusion (higher = more relevant, scale differs from cosine)"
+            ),
+            "results": results,
+        },
+        default=str,
+    )
+
+
+@mcp.tool()
+def get_file_skeleton(file_path: str) -> str:
+    """Get the API surface of a file (signatures, no bodies)."""
+    from coderay.skeleton.extractor import extract_skeleton
+
+    p = Path(file_path)
+    if not p.is_file():
+        return json.dumps({"error": f"File not found: {file_path}"})
+    content = p.read_text(encoding="utf-8", errors="replace")
+    return extract_skeleton(p, content)
+
+
+_STATIC_ANALYSIS_NOTE = (
+    "Based on static analysis of source code. Calls through dependency "
+    "injection, interfaces, dynamic dispatch (getattr), decorators, or "
+    "framework routing may not be detected."
+)
+
+
+@mcp.tool()
+def get_impact_radius(
+    node_id: str,
+    max_depth: int = 3,
+    index_dir: str | None = None,
+) -> str:
+    """Analyze the blast radius of changing a function or module."""
+    graph = _load_graph(index_dir)
+    if graph is None:
+        return json.dumps({"error": "No graph found. Run 'coderay build' first."})
+    impact = graph.get_impact_radius(node_id, depth=max_depth)
+    return json.dumps(
+        {
+            "results": [n.to_dict() for n in impact],
+            "note": _STATIC_ANALYSIS_NOTE,
+        }
+    )
+
+
+@mcp.tool()
+def index_status(index_dir: str | None = None) -> str:
+    """Check health and status of the semantic index."""
+    state = _load_state(index_dir)
+    if state is None:
+        return json.dumps({"status": "no_index", "message": "No index found."})
+
+    from coderay.core.config import get_embedding_dimensions, load_config
+    from coderay.state.version import read_index_version
+    from coderay.storage.lancedb import index_exists as idx_exists
+
+    idx_dir = _resolve_index_dir(index_dir)
+    has_store = idx_exists(idx_dir)
+    chunk_count = 0
+    if has_store:
+        from coderay.storage.lancedb import Store
+
+        config = load_config(idx_dir)
+        store = Store(idx_dir, dimensions=get_embedding_dimensions(config))
+        chunk_count = store.chunk_count()
+
+    return json.dumps(
+        {
+            "status": state.state.value,
+            "branch": state.branch,
+            "last_commit": state.last_commit,
+            "chunk_count": chunk_count,
+            "schema_version": read_index_version(idx_dir),
+            "has_store": has_store,
+        },
+        default=str,
+    )
+
+
+def main():
+    """Entry point for the coderay-mcp command."""
+    import sys
+
+    transport = "stdio"
+    if "--sse" in sys.argv:
+        transport = "sse"
+    mcp.run(transport=transport)
+
+
+if __name__ == "__main__":
+    main()

coderay/pipeline/indexer.py

@@ -0,0 +1,417 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from coderay.chunking.chunker import chunk_file
+from coderay.core.config import get_embedding_dimensions, load_config
+from coderay.core.timing import timed, timed_phase
+from coderay.core.utils import files_with_changed_content, hash_content, read_from_path
+from coderay.embedding.base import Embedder, load_embedder_from_config
+from coderay.graph.builder import build_and_save_graph
+from coderay.state.machine import IndexMeta, StateMachine
+from coderay.state.version import check_index_version, write_index_version
+from coderay.storage.lancedb import Store, index_exists
+from coderay.vcs.git import Git
+
+logger = logging.getLogger(__name__)
+
+RESUME_BATCH_SIZE = 200
+DEFAULT_REPO_ROOT = "."
+DEFAULT_INDEX_DIR = ".index"
+
+
+@dataclass
+class IndexResult:
+    """Result of an index build or update: cached, updated, and removed counts."""
+
+    cached: int = 0
+    updated: int = 0
+    removed: int = 0
+
+    def __str__(self) -> str:
+        return (
+            f"Cached: {self.cached}, Updated: {self.updated},"
+            f" Removed: {self.removed} chunks"
+        )
+
+
+class Indexer:
+    """Builds and maintains the semantic index for a repository."""
+
+    def __init__(
+        self,
+        repo_root: str | Path = DEFAULT_REPO_ROOT,
+        index_dir: str | Path = DEFAULT_INDEX_DIR,
+        config: dict[str, Any] | None = None,
+        embedder: Embedder | None = None,
+    ) -> None:
+        """Initialize the indexer."""
+        self._repo_root = Path(repo_root)
+        self._index_dir = Path(index_dir)
+        self._config = config or load_config(self._index_dir)
+        self._git = Git(self._repo_root)
+        self._state = StateMachine(self._index_dir)
+        self._embedder = embedder or load_embedder_from_config(self._config)
+        self._store = Store(
+            self._index_dir, dimensions=get_embedding_dimensions(self._config)
+        )
+        check_index_version(self._index_dir)
+
+    @property
+    def config(self) -> dict[str, Any]:
+        """Current config (embedder, index settings)."""
+        return self._config
+
+    @property
+    def repo_root(self) -> Path:
+        """Repository root path."""
+        return self._repo_root
+
+    @property
+    def index_dir(self) -> Path:
+        """Index directory path."""
+        return self._index_dir
+
+    @property
+    def current_state(self) -> IndexMeta | None:
+        """Current index meta state; None if no run has completed."""
+        return self._state.current_state
+
+    @timed("full_build")
+    def build_full(self) -> IndexResult:
+        """Full rebuild: discover, chunk, embed, and store all source files."""
+
+        current = self._state.current_state
+        last_branch = current.branch if current is not None else None
+        branch_switched = self._git.is_branch_switched(last_branch=last_branch)
+        if branch_switched:
+            return self.update_incremental()
+
+        current_run = current.current_run if current else None
+        saved_paths = current_run.paths_to_process if current_run else []
+        processed_count = current_run.processed_count if current_run else 0
+
+        can_resume = (
+            not branch_switched
+            and self._state.is_in_progress
+            and self._state.has_partial_progress
+        )
+
+        if can_resume:
+            paths_remaining = saved_paths[processed_count:]
+            if not paths_remaining:
+                self._state.finish(
+                    last_commit=self._git.get_head_commit(),
+                    branch=self._git.get_current_branch(),
+                )
+                write_index_version(self._index_dir)
+                self._refresh_graph()
+                return IndexResult(cached=len(self._state.file_hashes))
+            paths_to_process = paths_remaining
+            rel_paths = saved_paths
+        else:
+            self._state.set_incomplete()
+            self._store.clear()
+
+            py_files = self._git.discover_files()
+            if not py_files:
+                logger.warning("No source files found under %s", self._repo_root)
+                return IndexResult(cached=len(self._state.file_hashes))
+
+            rel_paths = [str(p.relative_to(self._repo_root)) for p in py_files]
+            paths_to_process = rel_paths
+            processed_count = 0
+            self._state.start(
+                branch=self._git.get_current_branch(),
+                last_commit=self._git.get_head_commit(),
+            )
+
+        all_path_hashes = self._run_batch_loop(
+            rel_paths=paths_to_process,
+            full_rel_paths=rel_paths,
+        )
+
+        if can_resume:
+            existing = self._state.file_hashes.copy()
+            existing.update(all_path_hashes)
+            self._state.file_hashes = existing
+        else:
+            self._state.file_hashes = all_path_hashes
+
+        self._state.finish(
+            last_commit=self._git.get_head_commit(),
+            branch=self._git.get_current_branch(),
+        )
+        write_index_version(self._index_dir)
+        self._refresh_graph()
+        return IndexResult(updated=len(self._state.file_hashes))
+
+    def update_incremental(self) -> IndexResult:
+        """Incremental update: re-index only changed, added, or deleted files."""
+
+        self._state.set_incomplete()
+
+        current = self._state.current_state
+        state_branch = current.branch if current else None
+        active_branch = self._git.get_current_branch()
+
+        if self._git.is_branch_switched(last_branch=state_branch):
+            logger.info(
+                "Branch switched %s -> %s; syncing index",
+                state_branch,
+                active_branch,
+            )
+            return self._sync_after_branch_switch()
+
+        to_add, to_remove = self._git.get_files_to_index(
+            last_commit=current.last_commit if current else None
+        )
+        if to_remove:
+            self._store.delete_by_paths(paths=to_remove)
+
+        # Remove deleted hashes from state
+        file_hashes = self._state.file_hashes.copy()
+        for path in to_remove:
+            file_hashes.pop(path, None)
+        # Idempotent if to_remove is empty
+        self._state.file_hashes = file_hashes
+
+        # Check what files are changed
+        changed_files = files_with_changed_content(
+            repo=self._repo_root, paths=to_add, file_hashes=file_hashes
+        )
+
+        if not changed_files and not to_remove:
+            self._state.finish(
+                last_commit=self._git.get_head_commit(),
+                branch=self._git.get_current_branch(),
+            )
+            self._refresh_graph()
+            logger.info("Nothing to update")
+            return IndexResult(cached=len(self._state.file_hashes))
+
+        return self._update(
+            paths_to_add=changed_files,
+            file_hashes=file_hashes,
+        )
+
+    def _sync_after_branch_switch(self) -> IndexResult:
+        """Sync index to current branch after a switch. Returns IndexResult."""
+        file_hashes = self._state.file_hashes.copy()
+        py_files = self._git.discover_files()
+
+        # All .py files were deleted from git
+        to_remove: list[str] = []
+        if not py_files:
+            to_remove = list(file_hashes)
+
+        if to_remove:
+            self._store.delete_by_paths(to_remove)
+            index_result = IndexResult(removed=len(file_hashes))
+            file_hashes.clear()
+
+            self._state.file_hashes = file_hashes
+            self._state.finish(
+                last_commit=self._git.get_head_commit(),
+                branch=self._git.get_current_branch(),
+            )
+            self._refresh_graph()
+            return index_result
+
+        rel_paths_current = {str(p.relative_to(self._repo_root)) for p in py_files}
+        # Deleted files on current branch
+        to_remove = [p for p in file_hashes if p not in rel_paths_current]
+        if to_remove:
+            self._store.delete_by_paths(to_remove)
+            for p in to_remove:
+                file_hashes.pop(p, None)
+
+        changed_files = files_with_changed_content(
+            repo=self._repo_root, paths=py_files, file_hashes=file_hashes
+        )
+
+        if not changed_files and not to_remove:
+            self._state.finish(
+                last_commit=self._git.get_head_commit(),
+                branch=self._git.get_current_branch(),
+            )
+            self._refresh_graph()
+            logger.info("Branch switch: index already in sync (no changes)")
+            return IndexResult(cached=len(self._state.file_hashes))
+
+        return self._update(
+            paths_to_add=changed_files,
+            file_hashes=file_hashes,
+        )
+
+    def _run_batch_loop(
+        self,
+        rel_paths: list[str],
+        full_rel_paths: list[str],
+    ) -> dict[str, str]:
+        """Run the pipeline in batches and save progress for resume."""
+        all_path_hashes: dict[str, str] = {}
+
+        for i in range(0, len(rel_paths), RESUME_BATCH_SIZE):
+            batch = rel_paths[i : i + RESUME_BATCH_SIZE]
+            all_path_hashes.update(self._run_pipeline(rel_paths=batch))
+            self._state.save_progress(
+                full_rel_paths=full_rel_paths,
+                processed_count=i + len(batch),
+            )
+
+        return all_path_hashes
+
+    def _update(
+        self,
+        paths_to_add: list[Path],
+        file_hashes: dict[str, str],
+    ) -> IndexResult:
+        """Run pipeline over paths, update hashes and state, then finish."""
+        rel_paths = [str(p.relative_to(self._repo_root)) for p in paths_to_add]
+        self._state.start(
+            branch=self._git.get_current_branch(),
+            last_commit=self._git.get_head_commit(),
+        )
+
+        batch_hashes = self._run_batch_loop(
+            rel_paths=rel_paths,
+            full_rel_paths=rel_paths,
+        )
+        file_hashes.update(batch_hashes)
+        self._state.file_hashes = file_hashes
+        self._state.finish(
+            last_commit=self._git.get_head_commit(),
+            branch=self._git.get_current_branch(),
+        )
+        self._refresh_graph(changed_paths=rel_paths)
+        return IndexResult(updated=len(batch_hashes))
+
+    @timed("pipeline")
+    def _run_pipeline(
+        self,
+        rel_paths: list[str],
+    ) -> dict[str, str]:
+        """Chunk, embed, and store the given files."""
+        files_content: list[tuple[str, str]] = []
+
+        with timed_phase("read"):
+            for p in rel_paths:
+                path = self._repo_root / p if not Path(p).is_absolute() else Path(p)
+                if not path.is_file():
+                    logger.warning("Skip (not a file): %s", p)
+                    continue
+                try:
+                    content = read_from_path(path)
+                    files_content.append((p, content))
+                except Exception as e:
+                    logger.warning("Skip (read failed) %s: %s", p, e)
+
+        if not files_content:
+            return {}
+
+        path_hashes = {p: hash_content(content) for p, content in files_content}
+
+        paths_to_replace = list({p for p, _ in files_content})
+        self._store.delete_by_paths(paths_to_replace)
+
+        with timed_phase("chunking"):
+            chunks = []
+            for p, content in files_content:
+                chunks.extend(chunk_file(p, content))
+
+        if not chunks:
+            logger.info("Pipeline done: 0 chunks in %d files", len(files_content))
+            return path_hashes
+
+        texts = [c.content for c in chunks]
+        with timed_phase("embedding"):
+            embeddings = self._embedder.embed(texts)
+
+        with timed_phase("storing"):
+            self._store.insert_chunks(chunks, embeddings)
+
+        logger.info(
+            "Pipeline done: %d chunks in %d files", len(chunks), len(files_content)
+        )
+        return path_hashes
+
+    def _refresh_graph(self, changed_paths: list[str] | None = None) -> None:
+        """Rebuild and save the code graph, logging warnings on failure."""
+        try:
+            build_and_save_graph(
+                self._repo_root,
+                self._index_dir,
+                changed_paths=changed_paths,
+            )
+        except Exception as e:
+            logger.warning("Graph refresh failed: %s", e)
+
+    def update_paths(
+        self,
+        changed: list[str],
+        removed: list[str] | None = None,
+    ) -> IndexResult:
+        """Update the index for explicit file paths (used by the file watcher)."""
+        self._state.set_incomplete()
+        file_hashes = self._state.file_hashes.copy()
+
+        removed = removed or []
+        if removed:
+            self._store.delete_by_paths(removed)
+            for p in removed:
+                file_hashes.pop(p, None)
+            self._state.file_hashes = file_hashes
+
+        if not changed:
+            self._state.finish(
+                last_commit=self._git.get_head_commit(),
+                branch=self._git.get_current_branch(),
+            )
+            if removed:
+                self._refresh_graph()
+            return IndexResult(removed=len(removed))
+
+        paths_to_add = [self._repo_root / p for p in changed]
+        existing = files_with_changed_content(
+            repo=self._repo_root, paths=paths_to_add, file_hashes=file_hashes
+        )
+
+        if not existing and not removed:
+            self._state.finish(
+                last_commit=self._git.get_head_commit(),
+                branch=self._git.get_current_branch(),
+            )
+            return IndexResult(cached=len(file_hashes))
+
+        if existing:
+            result = self._update(
+                paths_to_add=existing,
+                file_hashes=file_hashes,
+            )
+            result.removed = len(removed)
+            return result
+
+        self._state.finish(
+            last_commit=self._git.get_head_commit(),
+            branch=self._git.get_current_branch(),
+        )
+        self._refresh_graph()
+        return IndexResult(removed=len(removed))
+
+    def maintain(self) -> dict[str, Any]:
+        """Run store maintenance (compact and cleanup)."""
+        if not index_exists(self._index_dir):
+            return {}
+        return self._store.maintain()
+
+    def index_exists(self) -> bool:
+        """Return True if the index exists at index_dir."""
+        return index_exists(self._index_dir)
+
+    def error(self, exc: str) -> None:
+        """Mark the current run as errored with the given exception message."""
+        self._state.set_errored(exc=exc)