scripts.gen_codebase_map

Reference › scripts › gen_codebase_map

scripts/gen_codebase_map.py generates a compact, always-fresh map of the codebase for session ingestion. A SessionStart hook regenerates it, and a human can run it on demand. It imports orchestrata_ultimate.core.atomic_json for dump_text_atomic, and loads the serialize_codebase_ast toolkit by file path (_load_ast_toolkit) rather than by import name.

main takes plain flags from the command line: --with-full also writes the on-demand full symbol skeleton, --no-cache uses the uncached reference path, --cache PATH moves the cache from its default .cache/codebase-map.json, and --cache-stats prints hit statistics to standard error. build_map builds the header map, including the class map from _class_map (every class in src/ with its bases); build_full_map builds the full skeleton. It writes docs/CODEBASE_MAP.md, prints it to standard output, and with --with-full writes docs/CODEBASE_MAP_FULL.md through _write_full_map. _guard_and_write refuses to write the full map unless it holds at least 1000 signature lines, so a broken generation does not overwrite a good map.

_MapCache is a per-file cache keyed on file content. It stores the parse-derived data (symbol rows and first docstring line) so unchanged files are not re-parsed. All writes are atomic. The regeneration runs while the tree is being staged, and a truncating write keeps the inode, so git update-index hashing the file could die with SIGBUS; replacing the file avoids that.

At a glance

Class diagram

Dependency diagram

API

Classes

_MapCache

Line 241

Per-file cache of the two things the maps compute from a parse: the toolkit's symbol rows and the module's first docstring line.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/gen_codebase_map.py

#!/usr/bin/env python3
"""Generate a compact, always-fresh codebase map for session ingestion.

NOT the multi-megabyte serialized AST (``plans/ast-ultimate.json`` is ~8 MB —
tens of times a context window; a full ``symbols`` dump is ~770 KB). Both are far
too large to force into a session. This emits a size-bounded index — every runner
/ entrypoint script and every package, each with its one-line purpose — so the
*structure* of the tree is ambient in every session. The motivating failure: a
whole ThreadPool test runner (``scripts/run_tests_threadpool.py``) existed and was
never noticed because nothing put the inventory in front of the assistant.

Regenerated + injected at session start by
``.claude/hooks/inject-codebase-map.sh``. Detail (signatures, refs, callers) still
comes on demand from ``scripts/serialize_codebase_ast.py {symbols,refs,who-imports}``.
"""

from __future__ import annotations

import ast
import hashlib
import importlib.util
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "src" / "orchestrata_ultimate"
SCRIPTS = ROOT / "scripts"


def first_doc_line(path: Path) -> str:
    """First line of a module's docstring, or '' if none/unparseable."""
    try:
        mod = ast.parse(path.read_text(encoding="utf-8"))
    except OSError, SyntaxError, ValueError:
        return ""
    doc = ast.get_docstring(mod)
    if not doc:
        return ""
    return doc.strip().splitlines()[0].strip()


def rel(p: Path) -> str:
    return str(p.relative_to(ROOT))


def _section(
    add,
    title: str,
    paths: list[Path],
    *,
    with_counts: bool = False,
    cache: _MapCache | None = None,
) -> None:
    add(f"## {title} ({len(paths)})")
    for p in paths:
        target = p if p.is_file() else p / "__init__.py"
        doc = (cache.doc(target) if cache else first_doc_line(target)) or "—"
        if with_counts:
            nmods = len(list(p.rglob("*.py")))
            add(f"- `{p.name}/` ({nmods} modules) — {doc}")
        else:
            add(f"- `{rel(p)}` — {doc}")
    add("")


def _load_ast_toolkit():
    """Load ``serialize_codebase_ast`` BY FILE PATH, never by import name.

    ``import serialize_codebase_ast`` would require ``scripts/`` on
    ``sys.path``, and a module there whose name collides with a stdlib one
    then wins over it for the whole process — ``scripts/secrets.py`` shadowed
    stdlib ``secrets`` on 2026-08-12, broke ``redis.ConnectionPool``'s call to
    ``secrets.token_hex``, and took Redis + the gateway down so the suite could
    not start. File-path loading has no such failure mode.

    Returns the module, or ``None`` if it cannot be loaded — the class map is
    an enrichment, and a map without it beats no map at all.
    """
    # One load per process: --with-full reuses the toolkit _class_map loaded.
    target = SCRIPTS / "serialize_codebase_ast.py"
    cached = sys.modules.get("_gen_map_ast_toolkit")
    if cached is not None and getattr(cached, "__file__", None) == str(target):
        return cached
    try:
        spec = importlib.util.spec_from_file_location("_gen_map_ast_toolkit", target)
        if spec is None or spec.loader is None:
            return None
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)
        return module
    except (OSError, ImportError, SyntaxError) as exc:
        # NOT a silent swallow (Hard Rule 3): the reason is printed, and the
        # section itself renders as UNAVAILABLE so a degraded map is visible
        # rather than quietly incomplete. spec_from_file_location does NOT
        # raise on a missing file — exec_module does, with FileNotFoundError,
        # which the spec-is-None check never catches.
        print(
            f"gen_codebase_map: class map unavailable — {type(exc).__name__}: {exc}",
            file=sys.stderr,
        )
        return None


def _split_bases(signature: str) -> list[str]:
    """Base classes out of a ``class Foo(Bar, Generic[T])`` signature.

    Depth-aware: a naive ``split(",")`` would tear ``Mapping[str, int]`` into
    two bogus bases.
    """
    if "(" not in signature:
        return []
    inner = signature[signature.index("(") + 1 : signature.rindex(")")]
    bases: list[str] = []
    depth = 0
    current = ""
    for ch in inner:
        if ch in "[(":
            depth += 1
        elif ch in "])":
            depth -= 1
        if ch == "," and depth == 0:
            bases.append(current.strip())
            current = ""
            continue
        current += ch
    if current.strip():
        bases.append(current.strip())
    return [b for b in bases if b and "=" not in b]


def _base_key(base: str) -> str:
    """Bare class name for a base: ``pkg.mod.Base[T]`` -> ``Base``."""
    return base.split("[")[0].strip().rsplit(".", 1)[-1]


def _class_map(add, cache: _MapCache | None = None) -> None:
    """Every class in ``src/`` with its bases — the 'what do I extend' tier.

    WHY THIS EXISTS (measured 2026-08-12). The injected tier carried package
    purposes and a test inventory but NO class-level facts: of six abstractions
    needed to design one change correctly, five lived only in the NON-injected
    ``docs/CODEBASE_MAP_FULL.md``. The result was a third parallel
    "walk code and find things" implementation, written instead of extending
    ``orchestrata_ast.analyzer.CodebaseAnalyzer`` and the ``factory/detectors``
    rule engine that already existed.

    The old header/full split was cut at the wrong seam — by FILE rather than
    by SYMBOL KIND. Methods are 6,537 of 10,241 symbols and carry almost none
    of the extension signal; classes are 1,731 and carry nearly all of it. This
    section is ~35K tokens, about 3% of a 1M window, against ~404K for the full
    skeleton that the split was protecting subagents from.

    ``★N`` marks a class N other classes in this tree extend — the established
    extension points, and the answer to "what do I subclass".
    """
    toolkit = _load_ast_toolkit()
    if toolkit is None:
        add("## Class map — UNAVAILABLE (could not load the AST toolkit)")
        add("")
        return

    # ALL of src/, not just src/orchestrata_ultimate. orchestrata_ast lives
    # beside it and holds CodebaseAnalyzer — the very class whose absence from
    # context caused a parallel AST walker to be designed on 2026-08-12. A map
    # that omits the AST layer omits the thing most work needs to extend.
    src_root = ROOT / "src"
    if cache is not None:
        rows = _cached_list_symbols(toolkit, cache, "src", kind="class")
    else:
        rows = toolkit.list_symbols([str(src_root)], kind="class")

    # Subclass counts. Bases are matched on the BARE name because a base is
    # written as it was imported, not fully qualified — so this counts within
    # this tree only and is a navigation aid, not a resolved type graph. For a
    # resolved one use `code_intel.py` / `call-graph`, which resolve by type.
    local_names = {r["name"] for r in rows}
    subclass_count: dict[str, int] = {}
    for row in rows:
        for base in _split_bases(row["signature"]):
            key = _base_key(base)
            if key in local_names:
                subclass_count[key] = subclass_count.get(key, 0) + 1

    # list_symbols returns ABSOLUTE paths when called as an API and relative
    # ones from the CLI — normalise both against ROOT before slicing.
    by_pkg: dict[str, list[tuple[str, dict]]] = {}
    for row in rows:
        path = Path(row["file"])
        if not path.is_absolute():
            path = ROOT / path
        within = path.relative_to(src_root)
        # Group by top-level package, except inside orchestrata_ultimate where
        # that would collapse ~1600 classes into one heading — there, group by
        # its subpackage (agents, core, cxp, ...).
        parts = list(within.parts)
        if not parts:
            pkg = "(top level)"
        elif parts[0] == "orchestrata_ultimate" and len(parts) > 2:
            pkg = f"orchestrata_ultimate/{parts[1]}"
        else:
            pkg = parts[0]
        by_pkg.setdefault(pkg, []).append((rel(path), row))

    add(f"## Class map — every class in src/ ({len(rows)})")
    add("")
    add("> The tier that answers **what do I extend**, not merely what exists.")
    add("> Signatures carry base classes: `class Foo(Bar)` means Foo extends Bar.")
    add("> `★N` = N classes in this tree extend it — an established extension point.")
    add("> Bases match on bare name (as imported), so ★ counts are a navigation")
    add("> aid within this tree; for a TYPE-resolved graph use")
    add("> `scripts/code_intel.py` or `serialize_codebase_ast.py call-graph`.")
    add("> Methods are deliberately excluded — they are 64% of all symbols and")
    add("> carry almost none of the extension signal. For them, query on demand.")
    add("")

    for pkg in sorted(by_pkg):
        add(f"### {pkg} ({len(by_pkg[pkg])} classes)")
        for rel_path, row in by_pkg[pkg]:
            star = subclass_count.get(row["name"], 0)
            mark = f" ★{star}" if star else ""
            add(f"- `{rel_path}:{row['line']}` `{row['signature']}`{mark}")
        add("")


FULL_TREES = (
    "src/orchestrata_ultimate",
    "scripts",
    "rtm",
    "end2end",
    "tools",
    "harness",
    "architecture",
    "features",
    "specs",
    "claude",
)


class _MapCache:
    """Per-file cache of the two things the maps compute from a parse: the
    toolkit's symbol rows and the module's first docstring line.

    WHY THIS EXISTS (measured 2026-10-22). A ``--with-full`` run spent ~12 s
    re-parsing and re-walking ~2,300 files that had not changed since the last
    run (cProfile: ``_extract_symbols`` 12.8 s, ``first_doc_line`` 4.0 s). The
    toolkit's own warm cache is a 200 MB pickle that costs ~9 s just to load, so
    this keeps a small JSON of its own, keyed by file content (sha1), and
    re-parses only what changed.

    The ``version`` covers this file, the toolkit file and the interpreter: edit
    either script and the next run is one cold build, never a stale map. A
    missing, unreadable or malformed cache file starts empty — damage degrades
    to a cold build, it never breaks the map.
    """

    def __init__(self, path: Path, toolkit_file: Path) -> None:
        self.path = path
        digest = hashlib.sha1()
        for part in (Path(__file__), toolkit_file):
            try:
                digest.update(part.read_bytes())
            except OSError:
                pass
        digest.update(sys.version.encode("utf-8"))
        self._version = digest.hexdigest()
        self.hits = 0
        self.parsed = 0
        self._dirty = False
        self._files = self._load()
        self._used: dict[str, dict] = {}  # entries resolved this run
        self._done: set[tuple[str, str]] = set()  # (path, field) resolved this run
        # Trees kept between the docstring pass and the later symbol pass for
        # files both passes need, so such a file is parsed once per run.
        self._trees: dict[str, ast.Module] = {}

    def _load(self) -> dict[str, dict]:
        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
        except OSError, ValueError:
            return {}
        if not isinstance(data, dict) or data.get("version") != self._version:
            self._dirty = isinstance(data, dict)
            return {}
        files = data.get("files")
        return files if isinstance(files, dict) else {}

    def save(self, dump_text_atomic) -> None:
        """Persist when anything changed, keeping only files used this run."""
        if not self._dirty and set(self._files) == set(self._used):
            return
        self.path.parent.mkdir(parents=True, exist_ok=True)
        payload = {"version": self._version, "files": self._used}
        dump_text_atomic(self.path, json.dumps(payload, separators=(",", ":")))

    def stats(self) -> str:
        return (
            f"gen_codebase_map: cache: {self.hits} hit, {self.parsed} parsed, "
            f"{len(self._used)} files, {self.path}"
        )

    def _entry(self, path: Path) -> tuple[dict | None, str | None]:
        """The entry for ``path`` (current content), and its text when freshly
        read; ``(None, None)`` when the file cannot be read."""
        key = str(path)
        entry = self._used.get(key)
        if entry is not None:
            return entry, None
        try:
            text = path.read_text(encoding="utf-8")
        except OSError, UnicodeDecodeError:
            return None, None
        sha1 = hashlib.sha1(text.encode("utf-8")).hexdigest()
        old = self._files.get(key)
        if isinstance(old, dict) and old.get("sha1") == sha1:
            entry = old
        else:
            entry = {"sha1": sha1}
            self._dirty = True
        self._used[key] = entry
        return entry, text

    def _field(self, path: Path, name: str, default, compute):
        entry, text = self._entry(path)
        if entry is None:
            return default
        marker = (str(path), name)
        if name in entry:
            if marker not in self._done:
                self._done.add(marker)
                self.hits += 1
            return entry[name]
        if text is None:
            text = path.read_text(encoding="utf-8")
        entry[name] = compute(text)
        self._done.add(marker)
        self._dirty = True
        return entry[name]

    def _parse(self, path: Path, text: str, *, keep: bool) -> ast.Module | None:
        key = str(path)
        tree = self._trees.pop(key, None)
        if tree is None:
            self.parsed += 1
            try:
                tree = ast.parse(text)
            except SyntaxError, ValueError:
                return None
        if keep:
            self._trees[key] = tree
        return tree

    def symbols(self, path: Path, toolkit) -> list[dict] | None:
        """The toolkit's symbol rows for ``path`` with a root-relative ``file``,
        or None where the toolkit would skip the file."""

        def compute(text: str) -> list[dict] | None:
            tree = self._parse(path, text, keep=False)
            if tree is None:
                return None
            return toolkit._extract_symbols(tree, path.relative_to(ROOT))

        return self._field(path, "symbols", None, compute)

    def doc(self, path: Path) -> str:
        """Exactly ``first_doc_line(path)``."""

        def compute(text: str) -> str:
            # A file under src/ or a full-map tree is walked for symbols later
            # in the run, so its tree is kept for that second use.
            keep = any(
                path.is_relative_to(ROOT / tree) for tree in ("src", *FULL_TREES)
            )
            tree = self._parse(path, text, keep=keep)
            doc = ast.get_docstring(tree) if tree is not None else None
            return doc.strip().splitlines()[0].strip() if doc else ""

        return self._field(path, "doc", "", compute)


def _cached_list_symbols(toolkit, cache: _MapCache, source: str, kind=None):
    """``toolkit.list_symbols([ROOT / source], kind=kind)`` rebuilt from the
    cache — same files, same skips, same order, root-relative ``file`` values."""
    src = ROOT / source
    paths = [src] if src.is_file() else sorted(src.rglob("*.py"))
    seen: set[Path] = set()
    rows: list[dict] = []
    for path in paths:
        if any(skip in path.parts for skip in toolkit._DEP_SKIP_DIRS):
            continue
        resolved = path.resolve()
        if resolved in seen:
            continue
        seen.add(resolved)
        rows.extend(cache.symbols(path, toolkit) or [])
    if kind:
        rows = [r for r in rows if r["kind"] == kind]
    return sorted(rows, key=lambda r: (r["file"], r["line"]))


def build_full_map(toolkit, cache: _MapCache | None = None) -> str:
    """The on-demand full symbol skeleton, byte-identical to the hook's loop of
    ten ``serialize_codebase_ast.py symbols --source <tree>`` CLI calls — but
    from the one in-process toolkit load, so no warm-cache unpickle per tree."""
    # Anchors the relative tree strings so the output does not depend on cwd.
    toolkit.configure_repo_root(ROOT)
    lines = [
        "# Full symbol skeleton — src/ + scripts/ (class / function / method signatures)",
        "",
        "> ON-DEMAND tier (auto-generated with docs/CODEBASE_MAP.md; NOT @imported).",
        "> Read in ranges (grep -n first) — never whole.",
        "",
        "```",
    ]
    for tree in FULL_TREES:
        if not (ROOT / tree).is_dir():
            continue
        if cache is not None:
            rows = _cached_list_symbols(toolkit, cache, tree)
        else:
            rows = toolkit.list_symbols([tree])
        lines.append(f"### symbols: {tree}/")
        lines.extend(f"{r['file']}:{r['line']}: {r['signature']}" for r in rows)
        lines.append("")
        lines.append(f"{len(rows)} symbol(s).")
    lines.append("```")
    return "\n".join(lines) + "\n"


def build_map(cache: _MapCache | None = None) -> str:
    lines: list[str] = []
    add = lines.append

    add("# Codebase map (auto-generated at session start — do not hand-edit)")
    add("")
    add("> Compact index of runners, entrypoints, packages, and EVERY CLASS with its")
    add("> bases, derived from the live tree, so structure is always in context.")
    add("> The full AST (~8 MB) and the full symbol skeleton (~404K tokens, mostly")
    add("> methods) still do not fit — this does, at roughly 3% of a 1M window.")
    add("> For methods / references / callers, query the toolkit on demand:")
    add("> `scripts/serialize_codebase_ast.py {symbols,refs,who-imports,source}`.")
    add("")

    # 1. Runner / entrypoint scripts — the layer whose absence caused the miss.
    _section(
        add,
        "Runner & entrypoint scripts — scripts/",
        sorted(SCRIPTS.glob("*.py")),
        cache=cache,
    )

    # 2. src entrypoints: cli / main / server / app modules.
    seen: set[Path] = set()
    entries: list[Path] = []
    for g in ("**/cli*.py", "**/*_main.py", "**/main.py", "**/server.py", "**/app.py"):
        for p in sorted(SRC.glob(g)):
            if p not in seen:
                seen.add(p)
                entries.append(p)
    _section(
        add,
        "src entrypoints — cli / main / server / app",
        sorted(entries),
        cache=cache,
    )

    # 3. Package map.
    pkgs = sorted(
        d for d in SRC.iterdir() if d.is_dir() and (d / "__init__.py").exists()
    )
    _section(
        add,
        "Packages — src/orchestrata_ultimate/",
        pkgs,
        with_counts=True,
        cache=cache,
    )

    # 3b. Class map — the "what do I extend" tier. See _class_map's docstring
    #     for the measured failure that put it here.
    _class_map(add, cache)

    # 4. Test-module inventory (purpose only). Full test symbols are ~633K tokens
    #    — too large to always-load; this file-level inventory still makes every
    #    test's existence + purpose visible. For per-test detail, query on demand:
    #    scripts/serialize_codebase_ast.py symbols --source tests --pattern <name>
    tests = sorted((ROOT / "tests").rglob("test_*.py"))
    _section(
        add,
        "Test modules — tests/ (inventory; purpose only)",
        tests,
        cache=cache,
    )

    return "\n".join(lines) + "\n"


def main() -> int:
    # Atomic, not truncate-in-place: this map is regenerated by a SessionStart
    # hook while the tree it describes is being staged, and `git update-index`
    # mmaps a file to hash it. A truncating write keeps the inode, so the
    # reader faults past the new end and dies with SIGBUS -- measured twice on
    # 2026-08-19. See core.atomic_json.dump_text_atomic.
    from orchestrata_ultimate.core.atomic_json import dump_text_atomic

    args = sys.argv[1:]
    with_full = "--with-full" in args
    # One cache serves the header map, the class map and the full map, and is
    # saved once at the end. --no-cache runs the uncached reference path.
    cache: _MapCache | None = None
    if "--no-cache" not in args:
        cache_path = ROOT / ".cache" / "codebase-map.json"
        if "--cache" in args and args.index("--cache") + 1 < len(args):
            cache_path = Path(args[args.index("--cache") + 1])
        cache = _MapCache(cache_path, SCRIPTS / "serialize_codebase_ast.py")
    out = build_map(cache)
    dump_text_atomic(ROOT / "docs" / "CODEBASE_MAP.md", out)
    sys.stdout.write(out)
    if with_full:
        _write_full_map(dump_text_atomic, cache)
    if cache is not None:
        cache.save(dump_text_atomic)
        if "--cache-stats" in args:
            print(cache.stats(), file=sys.stderr)
    return 0


def _write_full_map(dump_text_atomic, cache: _MapCache | None = None) -> None:
    """Write docs/CODEBASE_MAP_FULL.md unless generation looks broken, in which
    case keep the existing file (a stale skeleton is recoverable, an empty one
    is not) and say so on stderr — the same guard the hook applied."""
    toolkit = _load_ast_toolkit()
    text = build_full_map(toolkit, cache) if toolkit is not None else ""
    _guard_and_write(text, ROOT / "docs" / "CODEBASE_MAP_FULL.md", dump_text_atomic)


def _guard_and_write(text: str, dest: Path, dump_text_atomic) -> int:
    """Write ``text`` to ``dest`` when it holds >= 1000 signature lines, else
    warn on stderr and keep ``dest``. Returns the signature count."""
    full = dest
    count = sum(1 for ln in text.splitlines() if ": def " in ln or ": class " in ln)
    if count >= 1000:
        dump_text_atomic(full, text)
    elif full.exists() and full.stat().st_size > 0:
        print(
            f"[codebase-map] WARNING: symbol generation produced {count} signatures; "
            "KEPT the existing docs/CODEBASE_MAP_FULL.md. Re-run: "
            "bash .claude/hooks/inject-codebase-map.sh",
            file=sys.stderr,
        )
    else:
        print(
            f"[codebase-map] WARNING: symbol generation produced {count} signatures "
            "and there is no existing docs/CODEBASE_MAP_FULL.md. The on-demand tier "
            "is UNAVAILABLE — use scripts/serialize_codebase_ast.py directly.",
            file=sys.stderr,
        )
    return count


if __name__ == "__main__":
    raise SystemExit(main())