heimdall-code 0.1.0__py3-none-any.whl

heimdall/mcp_server.py

@@ -0,0 +1,157 @@
+# Copyright 2026 The Heimdall Authors
+# SPDX-License-Identifier: Apache-2.0
+"""MCP server. One process serves every repository in the registry.
+
+Tools return precomputed relational answers -- callers grouped by depth, a resolved
+path, the source of a symbol -- not raw edges for the model to re-assemble.
+"""
+
+from __future__ import annotations
+
+import json
+from contextlib import closing as _closing
+from pathlib import Path
+from typing import Any
+
+try:  # the SDK renamed FastMCP to MCPServer; the decorator API is identical
+    from mcp.server.mcpserver import MCPServer as _Server
+except ImportError:  # pragma: no cover - older mcp releases
+    from mcp.server.fastmcp import FastMCP as _Server
+
+from . import registry
+from .graph import Graph
+
+mcp = _Server("heimdall")
+_DEFAULT: str | None = None
+
+
+def _open(project: str | None) -> Graph:
+    return Graph.open(registry.resolve(project or _DEFAULT), readonly=True)
+
+
+def _json(data: Any) -> str:
+    # compact: indentation is pure token cost to the model reading this
+    return json.dumps(data, separators=(",", ":"), default=str)
+
+
+@mcp.tool()
+def search(query: str, kind: str = "", limit: int = 20, project: str = "") -> str:
+    """Find symbols by name or signature, BM25 ranked. Start here when you only know
+    roughly what you are looking for.
+
+    kind: optional filter -- Function, Method, Class, Interface or File.
+    """
+    with _closing(_open(project)) as g:
+        return _json(g.search(query, limit, kind or None))
+
+
+@mcp.tool()
+def context(symbol: str, project: str = "") -> str:
+    """Everything about one symbol: where it is defined, its members, what it calls, who
+    calls it, and its supertypes. `symbol` may be a bare name, a qualified name
+    (`path/file.py:Class.method`) or a file path.
+    """
+    with _closing(_open(project)) as g:
+        return _json(g.context(symbol))
+
+
+@mcp.tool()
+def impact(symbol: str, depth: int = 3, min_confidence: float = 0.5, per_depth: int = 40, project: str = "") -> str:
+    """Blast radius: every symbol that transitively reaches this one, grouped by call
+    depth, plus the list of files affected. Use before changing or deleting anything.
+
+    min_confidence: call edges are scored 0.3 (name match only) to 1.0 (certain);
+    raise it to drop speculative edges.
+    """
+    with _closing(_open(project)) as g:
+        return _json(g.impact(symbol, depth, min_confidence, per_depth))
+
+
+@mcp.tool()
+def trace(source: str, target: str, max_depth: int = 8, min_confidence: float = 0.5, project: str = "") -> str:
+    """Shortest call/import/inheritance path from one symbol to another, or a statement
+    that none exists. Use to answer "how does A end up reaching B?".
+    """
+    with _closing(_open(project)) as g:
+        return _json(g.trace(source, target, max_depth, min_confidence))
+
+
+@mcp.tool()
+def overview(project: str = "") -> str:
+    """Shape of the repository: node and edge counts, languages, the most-imported files
+    and the most-called symbols. Use this first when you do not know the codebase.
+    """
+    with _closing(_open(project)) as g:
+        data = g.stats()
+        data["most_imported_files"] = g.query(
+            "SELECT n.path, count(*) AS importers FROM edge e JOIN node n ON n.id = e.dst"
+            " WHERE e.kind = 'IMPORTS' GROUP BY n.id ORDER BY importers DESC LIMIT 15"
+        )
+        data["most_called_symbols"] = g.query(
+            "SELECT n.qname, n.path, n.line, count(*) AS callers FROM edge e"
+            " JOIN node n ON n.id = e.dst WHERE e.kind = 'CALLS' AND e.confidence >= 0.7"
+            " GROUP BY n.id ORDER BY callers DESC LIMIT 15"
+        )
+        data["entry_points"] = g.query(
+            "SELECT n.qname, n.path, n.line FROM node n WHERE n.kind IN ('Function','Method')"
+            " AND n.name IN ('main','execute','run','handler','handle','serve','start')"
+            " ORDER BY n.path LIMIT 25"
+        )
+        return _json(data)
+
+
+@mcp.tool()
+def read_symbol(symbol: str, project: str = "") -> str:
+    """The actual source code of a symbol, read from disk. Cheaper than reading the whole
+    file when you only need one function or class.
+    """
+    with _closing(_open(project)) as g:
+        matches = g.resolve(symbol)
+        if not matches:
+            return _json({"error": f"no symbol matching {symbol!r}", "suggestions": g.search(symbol, 5)})
+        n = matches[0]
+        root = Path(g.meta().get("root", "."))
+        try:
+            lines = (root / n["path"]).read_text(encoding="utf-8", errors="replace").splitlines()
+        except OSError as exc:
+            return _json({"error": str(exc)})
+        start, end = max(0, n["line"] - 1), (n["end_line"] or n["line"]) or len(lines)
+        body = lines[start:end] if n["kind"] not in ("File", "Folder") else lines
+        return _json(
+            {
+                "qname": n["qname"],
+                "path": n["path"],
+                "lines": f"{start + 1}-{start + len(body)}",
+                "source": "\n".join(body),
+            }
+        )
+
+
+@mcp.tool()
+def sql(query: str, limit: int = 100, project: str = "") -> str:
+    """Read-only SQL against the graph, for questions the other tools do not cover.
+
+    Schema:
+      node(id, kind, name, qname, path, line, end_line, lang, signature)
+        kind: Folder | File | Class | Interface | Function | Method
+      edge(src, dst, kind, confidence, line)
+        kind: CONTAINS | DEFINES | IMPORTS | CALLS | EXTENDS
+    Recursive CTEs work; only a single SELECT or WITH statement is accepted.
+    """
+    with _closing(_open(project)) as g:
+        try:
+            return _json(g.query(query, limit))
+        except Exception as exc:
+            return _json({"error": f"{type(exc).__name__}: {exc}"})
+
+
+@mcp.tool()
+def projects() -> str:
+    """Repositories that have been indexed and can be passed as `project`."""
+    return _json(registry.load())
+
+
+def serve(project: str | None = None) -> None:
+    global _DEFAULT
+    _DEFAULT = project
+    mcp.run()

heimdall/parse.py

@@ -0,0 +1,409 @@
+# Copyright 2026 The Heimdall Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Tree-sitter symbol extraction.
+
+Language support is a data table, not code: node-type -> symbol kind. Tree-sitter
+grammars name their nodes consistently enough that one generic walker covers every
+language here. Adding a language means adding rows, not a module.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from tree_sitter_language_pack import get_parser
+
+# ---------------------------------------------------------------- language data
+
+EXT_LANG = {
+    ".py": "python",
+    ".pyi": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".mjs": "javascript",
+    ".cjs": "javascript",
+    ".ts": "typescript",
+    ".tsx": "tsx",
+    ".go": "go",
+    ".rs": "rust",
+    ".java": "java",
+    ".kt": "kotlin",
+    ".kts": "kotlin",
+    ".rb": "ruby",
+    ".php": "php",
+    ".cs": "csharp",
+    ".swift": "swift",
+    ".scala": "scala",
+    ".c": "c",
+    ".h": "c",
+    ".cc": "cpp",
+    ".cpp": "cpp",
+    ".cxx": "cpp",
+    ".hpp": "cpp",
+    ".lua": "lua",
+    ".dart": "dart",
+    ".sh": "bash",
+}
+
+_JS_DEFS = {
+    "function_declaration": "Function",
+    "generator_function_declaration": "Function",
+    "class_declaration": "Class",
+    "method_definition": "Method",
+}
+_TS_DEFS = _JS_DEFS | {
+    "interface_declaration": "Interface",
+    "abstract_class_declaration": "Class",
+    "enum_declaration": "Class",
+    "type_alias_declaration": "Interface",
+    "function_signature": "Function",
+}
+
+# node type -> symbol kind. A Function nested directly in a Class becomes a Method.
+DEFS: dict[str, dict[str, str]] = {
+    "python": {"function_definition": "Function", "class_definition": "Class"},
+    "javascript": _JS_DEFS,
+    "typescript": _TS_DEFS,
+    "tsx": _TS_DEFS,
+    "go": {"function_declaration": "Function", "method_declaration": "Method", "type_spec": "Class"},
+    "rust": {
+        "function_item": "Function",
+        "struct_item": "Class",
+        "trait_item": "Interface",
+        "enum_item": "Class",
+        "impl_item": "Class",
+    },
+    "java": {
+        "class_declaration": "Class",
+        "interface_declaration": "Interface",
+        "method_declaration": "Method",
+        "constructor_declaration": "Method",
+        "enum_declaration": "Class",
+        "record_declaration": "Class",
+    },
+    "kotlin": {"class_declaration": "Class", "function_declaration": "Function", "object_declaration": "Class"},
+    "csharp": {
+        "class_declaration": "Class",
+        "interface_declaration": "Interface",
+        "method_declaration": "Method",
+        "struct_declaration": "Class",
+        "record_declaration": "Class",
+        "constructor_declaration": "Method",
+    },
+    "ruby": {"class": "Class", "module": "Class", "method": "Method", "singleton_method": "Method"},
+    "php": {
+        "class_declaration": "Class",
+        "interface_declaration": "Interface",
+        "function_definition": "Function",
+        "method_declaration": "Method",
+        "trait_declaration": "Interface",
+    },
+    "c": {"function_definition": "Function", "struct_specifier": "Class"},
+    "cpp": {
+        "function_definition": "Function",
+        "class_specifier": "Class",
+        "struct_specifier": "Class",
+        "namespace_definition": "Class",
+    },
+    "swift": {"class_declaration": "Class", "function_declaration": "Function", "protocol_declaration": "Interface"},
+    "scala": {
+        "class_definition": "Class",
+        "object_definition": "Class",
+        "trait_definition": "Interface",
+        "function_definition": "Function",
+    },
+    "lua": {"function_declaration": "Function"},
+    "dart": {"class_definition": "Class", "function_signature": "Function", "method_signature": "Method"},
+    "bash": {"function_definition": "Function"},
+}
+
+CALL_NODES: dict[str, set[str]] = {
+    "python": {"call"},
+    "javascript": {"call_expression", "new_expression"},
+    "typescript": {"call_expression", "new_expression"},
+    "tsx": {"call_expression", "new_expression"},
+    "go": {"call_expression"},
+    "rust": {"call_expression", "macro_invocation"},
+    "java": {"method_invocation", "object_creation_expression"},
+    "kotlin": {"call_expression"},
+    "csharp": {"invocation_expression", "object_creation_expression"},
+    "ruby": {"call", "method_call"},
+    "php": {"function_call_expression", "member_call_expression", "scoped_call_expression"},
+    "c": {"call_expression"},
+    "cpp": {"call_expression"},
+    "swift": {"call_expression"},
+    "scala": {"call_expression"},
+    "lua": {"function_call"},
+    "dart": {"selector"},
+    "bash": {"command"},
+}
+
+IMPORT_NODES: dict[str, set[str]] = {
+    "python": {"import_statement", "import_from_statement"},
+    "javascript": {"import_statement", "export_statement"},
+    "typescript": {"import_statement", "export_statement"},
+    "tsx": {"import_statement", "export_statement"},
+    "go": {"import_spec"},
+    "rust": {"use_declaration"},
+    "java": {"import_declaration"},
+    "kotlin": {"import_header"},
+    "csharp": {"using_directive"},
+    "ruby": set(),
+    "php": {"namespace_use_declaration"},
+    "c": {"preproc_include"},
+    "cpp": {"preproc_include"},
+    "swift": {"import_declaration"},
+    "scala": {"import_declaration"},
+    "lua": set(),
+    "dart": {"configurable_uri"},
+    "bash": set(),
+}
+
+# where a base class / implemented interface hides: a named field on some grammars,
+# a plain child node on others
+HERITAGE_FIELDS = ("superclasses", "superclass", "super_interfaces", "interfaces")
+HERITAGE_TYPES = {
+    "class_heritage",
+    "superclass",
+    "super_interfaces",
+    "extends_clause",
+    "implements_clause",
+    "base_class_clause",
+    "interface_type_list",
+    "base_list",
+    "delegation_specifier",
+    "extends_type_clause",
+    "inheritance_specifier",
+}
+
+# `from x import y as z` / `import { y as z }`: the call site says `z`, the symbol is `y`.
+# Both grammars spell it with `name` and `alias` fields, so one set covers both.
+# ponytail: Python and the JS family, where renaming on import is idiomatic. Go and
+# Kotlin spell it differently; add their node types here if their edges start to matter.
+ALIAS_NODES = {"aliased_import", "import_specifier"}
+
+_NAME_FIELDS = ("name", "declarator", "type", "path")
+_MODULE_FIELDS = ("module_name", "source", "path", "name")
+_KEYWORDS = {"import", "from", "use", "using", "require", "include", "export", "package", ";", "#include"}
+
+
+# ---------------------------------------------------------------- data classes
+
+
+@dataclass(slots=True)
+class Symbol:
+    kind: str
+    name: str
+    qname: str
+    line: int
+    end_line: int
+    parent: str | None
+    signature: str
+
+
+@dataclass(slots=True)
+class FileParse:
+    path: str
+    lang: str
+    symbols: list[Symbol] = field(default_factory=list)
+    # (caller_qname_or_None, callee_name, line, receiver_or_None)
+    calls: list[tuple[str | None, str, int, str | None]] = field(default_factory=list)
+    # (module, line, enclosing_symbol_qname_or_None) -- the scope matters: a function
+    # that does its own `from x import y` must not be resolved against the whole file
+    imports: list[tuple[str, int, str | None]] = field(default_factory=list)
+    # (local_name, original_name, line_of_its_import) for renamed imports
+    renames: list[tuple[str, str, int]] = field(default_factory=list)
+    # (class_qname, base_name)
+    heritage: list[tuple[str, str]] = field(default_factory=list)
+    error: str | None = None
+
+
+# ---------------------------------------------------------------- helpers
+
+
+def _text(node) -> str:
+    return node.text.decode("utf-8", "replace")
+
+
+def _name_of(node) -> str | None:
+    for f in _NAME_FIELDS:
+        n = node.child_by_field_name(f)
+        if n is not None:
+            t = _text(n)
+            # C declarators carry the params; keep the identifier only.
+            return t.split("(")[0].strip("*& ") or None
+    # rust `impl Foo`, ruby `class Foo` and friends expose the name positionally.
+    for c in node.children:
+        if c.type in ("identifier", "type_identifier", "constant", "field_identifier", "name"):
+            return _text(c)
+    return None
+
+
+def _callee_of(node) -> tuple[str | None, str | None]:
+    """(receiver, name) for a call. `self.n()` -> ("self", "n"), `foo()` -> (None, "foo")."""
+    if node.type == "selector":
+        # dart puts the callee and its argument list side by side, not nested
+        if not any(c.type == "argument_part" for c in node.children):
+            return None, None
+        prev = node.prev_named_sibling
+        return (None, _text(prev).lstrip(".")) if prev is not None else (None, None)
+    n = node.child_by_field_name("function") or node.child_by_field_name("name") or node.child_by_field_name("method")
+    if n is None:
+        for c in node.children:
+            if c.type in (
+                "identifier",
+                "simple_identifier",
+                "field_expression",
+                "selector_expression",
+                "member_expression",
+                "navigation_expression",
+                "attribute",
+                "scoped_identifier",
+            ):
+                n = c
+                break
+    if n is None:
+        return None, None
+    parts = _text(n).replace("::", ".").split(".")
+    name = parts[-1].strip()
+    receiver = parts[-2].strip() if len(parts) > 1 else None
+    return receiver or None, name or None
+
+
+def _module_of(node) -> str | None:
+    # `export class Foo ...` is not an import; only `export ... from "x"` is.
+    if node.type.endswith("export_statement"):
+        src = node.child_by_field_name("source")
+        return _text(src).strip("\"'<>") if src is not None else None
+    for f in _MODULE_FIELDS:
+        n = node.child_by_field_name(f)
+        if n is not None:
+            return _clean_module(_text(n))
+    for c in node.children:
+        if _text(c) in _KEYWORDS or not c.is_named:
+            continue
+        return _clean_module(_text(c))
+    return None
+
+
+def _renames(import_node) -> list[tuple[str, str]]:
+    """(local, original) for every `as` rename inside one import statement."""
+    found: list[tuple[str, str]] = []
+
+    def walk(n):
+        if n.type in ALIAS_NODES:
+            orig, local = n.child_by_field_name("name"), n.child_by_field_name("alias")
+            if orig is not None and local is not None:
+                found.append((_text(local).strip(), _text(orig).split(".")[-1].strip()))
+            return
+        for c in n.children:
+            walk(c)
+
+    walk(import_node)
+    return found
+
+
+def _clean_module(text: str) -> str | None:
+    """A module path is a single token; anything with whitespace is a mis-parse."""
+    text = text.strip().strip("\"'<>")
+    return text if text and not any(c.isspace() for c in text) else None
+
+
+def _signature(node, src: bytes) -> str:
+    """First line of the declaration, which is the signature for every language here."""
+    end = src.find(b"\n", node.start_byte)
+    end = node.end_byte if end == -1 else min(end, node.end_byte)
+    return src[node.start_byte : end].decode("utf-8", "replace").strip()[:400]
+
+
+# ---------------------------------------------------------------- the walker
+
+
+def parse_source(path: str, lang: str, src: bytes) -> FileParse:
+    out = FileParse(path=path, lang=lang)
+    defs = DEFS.get(lang, {})
+    calls = CALL_NODES.get(lang, set())
+    imports = IMPORT_NODES.get(lang, set())
+    try:
+        tree = get_parser(lang).parse(src)
+    except Exception as exc:  # unparsable file: record and keep indexing the repo
+        out.error = f"{type(exc).__name__}: {exc}"
+        return out
+
+    def visit(node, parent: Symbol | None) -> Symbol | None:
+        # Python hangs the decorators and the `def` off one node, decorators first. Visit
+        # the def so its Symbol exists, then the decorators against it: `@setup(...)` is a
+        # call by the function it wraps, not by the class that happens to enclose it.
+        # ponytail: other languages nest decorators *inside* the definition node, where
+        # the walk already gets the parent right.
+        if node.type == "decorated_definition":
+            inner = node.child_by_field_name("definition")
+            sym = visit(inner, parent) if inner is not None else None
+            for c in node.children:
+                if c.type == "decorator":
+                    visit(c, sym or parent)
+            return sym
+
+        kind = defs.get(node.type)
+        if kind is not None:
+            name = _name_of(node)
+            if name:
+                if kind == "Function" and parent is not None and parent.kind in ("Class", "Interface"):
+                    kind = "Method"
+                qname = f"{parent.qname}.{name}" if parent else f"{path}:{name}"
+                sym = Symbol(
+                    kind=kind,
+                    name=name,
+                    qname=qname,
+                    line=node.start_point[0] + 1,
+                    end_line=node.end_point[0] + 1,
+                    parent=parent.qname if parent else None,
+                    signature=_signature(node, src),
+                )
+                out.symbols.append(sym)
+                if kind in ("Class", "Interface"):
+                    clauses = [node.child_by_field_name(f) for f in HERITAGE_FIELDS]
+                    clauses += [c for c in node.children if c.type in HERITAGE_TYPES]
+                    bases: list[str] = []
+                    for h in clauses:
+                        if h is not None:
+                            bases += [b for b in _bases(h) if b not in bases]
+                    out.heritage += [(qname, b) for b in bases]
+                for c in node.children:
+                    visit(c, sym)
+                return sym
+        elif node.type in calls:
+            receiver, callee = _callee_of(node)
+            if callee and callee.isidentifier():
+                out.calls.append((parent.qname if parent else None, callee, node.start_point[0] + 1, receiver))
+        elif node.type in imports:
+            mod = _module_of(node)
+            if mod:
+                line = node.start_point[0] + 1
+                out.imports.append((mod, line, parent.qname if parent else None))
+                # keyed by the statement's line so the resolver can tell whether the
+                # module it came from is even in this repo
+                out.renames += [(local, orig, line) for local, orig in _renames(node)]
+        for c in node.children:
+            visit(c, parent)
+        return None
+
+    visit(tree.root_node, None)
+    return out
+
+
+def _bases(heritage_node) -> list[str]:
+    """Identifiers inside a heritage clause, e.g. `extends B implements I` -> [B, I]."""
+    found: list[str] = []
+
+    def walk(n):
+        if n.type in ("identifier", "type_identifier", "constant", "scoped_type_identifier", "generic_type"):
+            t = _text(n).split("<")[0].split(".")[-1].strip()
+            if t and t.isidentifier():
+                found.append(t)
+            return
+        for c in n.children:
+            walk(c)
+
+    walk(heritage_node)
+    return found

heimdall/registry.py

@@ -0,0 +1,51 @@
+# Copyright 2026 The Heimdall Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Which repositories have been indexed, so one MCP server can serve all of them."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+HOME = Path(os.environ.get("HEIMDALL_HOME") or Path.home() / ".heimdall")
+REGISTRY = HOME / "registry.json"
+
+
+def db_path_for(root: Path) -> Path:
+    return root.resolve() / ".heimdall" / "graph.db"
+
+
+def load() -> dict[str, str]:
+    try:
+        return json.loads(REGISTRY.read_text())
+    except (OSError, ValueError):
+        return {}
+
+
+def register(name: str, db: Path) -> None:
+    entries = load()
+    entries[name] = str(db.resolve())
+    REGISTRY.parent.mkdir(parents=True, exist_ok=True)
+    REGISTRY.write_text(json.dumps(entries, indent=2, sort_keys=True))
+
+
+def resolve(project: str | None, cwd: Path | None = None) -> Path:
+    """A project name, an explicit db path, or the repo in the current directory."""
+    if project:
+        p = Path(project)
+        if p.suffix == ".db" and p.exists():
+            return p
+        entries = load()
+        if project in entries:
+            return Path(entries[project])
+        matches = [v for k, v in entries.items() if project in k]
+        if len(matches) == 1:
+            return Path(matches[0])
+        raise SystemExit(f"unknown project {project!r}; known: {', '.join(sorted(load())) or '(none)'}")
+    here = (cwd or Path.cwd()).resolve()
+    for d in (here, *here.parents):
+        db = d / ".heimdall" / "graph.db"
+        if db.exists():
+            return db
+    raise SystemExit("no index found here -- run `heimdall index` first, or pass --project")