server.py

"""
codensa MCP Server

Tools exposed to AI coding assistants:
  project_summary      - Compressed project overview (replaces reading all files)
  semantic_search      - Find relevant code chunks by natural language
  file_graph           - Dependencies/callers for a specific file
  find_symbol          - Locate where a class/function is defined
  index_status         - Check if index is current
  reindex              - Force a fresh index run
  get_file_content     - Read a file (with token budget guard)
  list_files           - Browse the file tree
"""

import os
import sys
import json
from pathlib import Path
from typing import Optional

import fastmcp
from fastmcp import FastMCP

# ── Resolve project root ───────────────────────────────────────────────────────
# Priority: PROJECT_ROOT env var → cwd → first CLI arg

def resolve_project_root() -> str:
    if "PROJECT_ROOT" in os.environ:
        return os.environ["PROJECT_ROOT"]
    if len(sys.argv) > 1 and Path(sys.argv[1]).is_dir():
        return sys.argv[1]
    return os.getcwd()


PROJECT_ROOT = resolve_project_root()

# Lazy-import indexer so the MCP server starts fast
from src.indexer import CodensaIndexer

_indexer: Optional[CodensaIndexer] = None


def get_indexer() -> CodensaIndexer:
    global _indexer
    if _indexer is None:
        _indexer = CodensaIndexer(PROJECT_ROOT)
        # Auto-index on first use
        _indexer.index(verbose=True)
        _indexer.load()
    return _indexer


# ── MCP Server ─────────────────────────────────────────────────────────────────

mcp = FastMCP(
    "codensa",
    instructions="""
You are connected to a GraphRAG knowledge store for the current codebase.
ALWAYS call project_summary first at the start of a session — it gives you a 
compressed map of the whole project so you do NOT need to read individual files 
unless you need exact content. Use semantic_search to find relevant code before 
asking for file contents. This saves tokens and gives you better context.
""",
)


@mcp.tool(
    description=(
        "Returns a compressed markdown summary of the entire project: "
        "tech stack, language distribution, file map, and graph stats. "
        "Call this FIRST at the start of every session to understand the project "
        "without reading individual files. Token-efficient."
    )
)
def project_summary() -> str:
    idx = get_indexer()
    return idx.get_summary()


@mcp.tool(
    description=(
        "Semantic search over all indexed code chunks. "
        "Use this to find relevant files/functions before opening any file. "
        "Returns the most relevant code snippets with file paths and relevance scores."
    )
)
def semantic_search(
    query: str,
    top_k: int = 6,
) -> str:
    """
    Args:
        query: Natural language question or description of code you need
        top_k: Number of results (1-10, default 6)
    """
    idx = get_indexer()
    top_k = max(1, min(top_k, 10))
    hits = idx.query(query, top_k=top_k)
    if not hits:
        return "No results found. Try broader search terms."
    lines = [f"## Semantic Search: {query!r}\n"]
    for i, hit in enumerate(hits, 1):
        lines.append(f"### [{i}] `{hit['file']}` (score: {hit['score']})")
        lines.append("```")
        lines.append(hit["text"][:800])
        lines.append("```\n")
    return "\n".join(lines)


@mcp.tool(
    description=(
        "Returns the dependency graph for a specific file: "
        "what it imports, which files import it, and what symbols it defines. "
        "Use to understand the blast radius of a change without reading file contents."
    )
)
def file_graph(filepath: str) -> str:
    """
    Args:
        filepath: Relative path from project root (e.g. 'src/api/routes.py')
    """
    idx = get_indexer()
    result = idx.get_file_graph(filepath)
    if not result:
        return f"File `{filepath}` not found in index. Run reindex if file is new."
    return json.dumps(result, indent=2)


@mcp.tool(
    description=(
        "Find which file(s) define a class, function, or interface by name. "
        "Faster than grep — uses the graph index. Returns file paths + symbol IDs."
    )
)
def find_symbol(symbol_name: str) -> str:
    """
    Args:
        symbol_name: Class or function name (partial match supported)
    """
    idx = get_indexer()
    results = idx.get_symbol(symbol_name)
    if not results:
        return f"Symbol `{symbol_name}` not found. Try partial name or reindex."
    lines = [f"## Symbol: `{symbol_name}`\n"]
    for r in results:
        lines.append(f"- `{r['symbol']}` → defined in `{r['file']}`")
    return "\n".join(lines)


@mcp.tool(
    description=(
        "List files in the project matching a pattern. "
        "Use to browse directory structure without shell access."
    )
)
def list_files(pattern: str = "", directory: str = "") -> str:
    """
    Args:
        pattern: Substring to filter filenames (e.g. 'api', '.test.ts', 'model')
        directory: Limit to files under this subdirectory
    """
    idx = get_indexer()
    files = idx.search_files(pattern or directory or "")
    if directory:
        files = [f for f in files if f.startswith(directory)]
    if not files:
        return "No matching files found in index."
    files.sort()
    lines = [f"## Files matching `{pattern or '*'}`\n"]
    for f in files[:100]:
        node = idx.graph.nodes.get(f, {})
        size = node.get("size_kb", "?")
        syms = ", ".join(node.get("symbols", [])[:4])
        lines.append(f"- `{f}` ({size}KB) {f'| {syms}' if syms else ''}")
    if len(files) > 100:
        lines.append(f"\n... and {len(files) - 100} more")
    return "\n".join(lines)


@mcp.tool(
    description=(
        "Read the raw content of a single file. "
        "Only call this when you need exact code — prefer semantic_search first. "
        "Enforces a token budget to prevent context overflow."
    )
)
def get_file_content(
    filepath: str,
    max_tokens: int = 3000,
    start_line: int = 1,
    end_line: Optional[int] = None,
) -> str:
    """
    Args:
        filepath: Relative path from project root
        max_tokens: Max tokens to return (default 3000, max 8000)
        start_line: First line to return (1-indexed)
        end_line: Last line to return (optional)
    """
    from src.indexer import count_tokens, truncate_to_tokens
    max_tokens = min(max_tokens, 8000)
    full_path = Path(PROJECT_ROOT) / filepath
    if not full_path.exists():
        return f"File not found: {filepath}"
    try:
        content = full_path.read_text(encoding="utf-8", errors="ignore")
        lines = content.splitlines()
        sl = max(0, start_line - 1)
        el = end_line if end_line else len(lines)
        slice_content = "\n".join(lines[sl:el])
        truncated = truncate_to_tokens(slice_content, max_tokens)
        token_count = count_tokens(truncated)
        return f"## `{filepath}` (lines {start_line}-{el}, ~{token_count} tokens)\n\n```\n{truncated}\n```"
    except Exception as e:
        return f"Error reading {filepath}: {e}"


@mcp.tool(
    description=(
        "Check the index status: when it was last built, how many files/chunks, "
        "and whether any files have changed since the last index run."
    )
)
def index_status() -> str:
    idx = get_indexer()
    meta = idx.meta
    if not meta:
        return "Project not yet indexed. Run reindex."
    import time
    ts = meta.get("indexed_at", 0)
    age_min = round((time.time() - ts) / 60, 1)
    stale = idx.needs_reindex()
    lines = [
        f"## Index Status",
        f"- **Root**: `{meta.get('project_root')}`",
        f"- **Indexed**: {time.strftime('%Y-%m-%d %H:%M', time.localtime(ts))} ({age_min} min ago)",
        f"- **Files**: {meta.get('file_count', 0)}",
        f"- **Chunks**: {meta.get('chunk_count', 0)}",
        f"- **Graph nodes**: {meta.get('node_count', 0)}",
        f"- **Graph edges**: {meta.get('edge_count', 0)}",
        f"- **Stale**: {'⚠️ YES — run reindex' if stale else '✅ No'}",
        f"- **Elapsed**: {meta.get('elapsed_s', '?')}s",
    ]
    tech = meta.get("stats", {}).get("project_types", [])
    if tech:
        lines.append(f"- **Tech**: {', '.join(tech)}")
    return "\n".join(lines)


@mcp.tool(
    description=(
        "Force a full re-index of the project. "
        "Run this after adding new files, renaming modules, or major refactors. "
        "Small edits are auto-detected on next query."
    )
)
def reindex(force: bool = True) -> str:
    global _indexer
    idx = get_indexer()
    meta = idx.index(force=force, verbose=True)
    idx.load()
    return (
        f"✅ Reindex complete: {meta.get('file_count')} files, "
        f"{meta.get('chunk_count')} chunks, "
        f"{meta.get('node_count')} graph nodes, "
        f"took {meta.get('elapsed_s')}s"
    )


@mcp.tool(
    description=(
        "Get the most-connected files in the dependency graph. "
        "High in-degree = heavily imported (core modules). "
        "Useful to understand project architecture quickly."
    )
)
def hot_files(top_n: int = 15) -> str:
    """
    Args:
        top_n: Number of top files to return (default 15)
    """
    idx = get_indexer()
    file_nodes = [(n, d) for n, d in idx.graph.nodes(data=True) if d.get("type") == "file"]
    if not file_nodes:
        return "No file nodes in graph. Run reindex."
    scored = []
    for n, d in file_nodes:
        in_deg = idx.graph.in_degree(n)
        out_deg = idx.graph.out_degree(n)
        scored.append((n, d, in_deg, out_deg))
    scored.sort(key=lambda x: x[2], reverse=True)
    lines = [f"## Top {top_n} Most-Imported Files\n"]
    for n, d, ind, outd in scored[:top_n]:
        syms = ", ".join(d.get("symbols", [])[:4])
        lines.append(
            f"- `{n}` | imported-by: **{ind}** | imports: {outd}"
            + (f" | `{syms}`" if syms else "")
        )
    return "\n".join(lines)


# ── Entry point ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="codensa MCP server")
    parser.add_argument("project_root", nargs="?", default=None,
                        help="Path to project root (default: cwd or PROJECT_ROOT env)")
    parser.add_argument("--transport", default="stdio", choices=["stdio", "streamable-http"],
                        help="MCP transport (default: stdio)")
    parser.add_argument("--port", type=int, default=8765,
                        help="Port for HTTP transport")
    parser.add_argument("--reindex", action="store_true",
                        help="Force reindex before starting")
    args = parser.parse_args()

    if args.project_root:
        os.environ["PROJECT_ROOT"] = args.project_root
        # Update module-level variable via globals()
        globals()["PROJECT_ROOT"] = args.project_root

    if args.reindex:
        idx = get_indexer()
        idx.index(force=True, verbose=True)

    print(f"[codensa] Starting MCP server for: {PROJECT_ROOT}", file=sys.stderr)
    print(f"[codensa] Transport: {args.transport}", file=sys.stderr)

    if args.transport == "stdio":
        mcp.run(transport="stdio")
    else:
        mcp.run(transport="streamable-http", port=args.port)