orchestrata_ultimate.tools.codebase_query_tools

Reference › orchestrata_ultimate › tools › codebase_query_tools

This module wires four of the Prolog codebase-intelligence queries in codebase_tools into the unified registry as BaseTool classes (SPEC-062 and SPEC-072). agents.tool_registry imports it, and tests/tools/test_codebase_query_tools.py covers it.

The four tools each call one CodebaseTools method: CodebaseImpactTool (everything that transitively depends on a symbol), CodebaseModuleDepsTool (transitive internal module dependencies), CodebaseFindCyclesTool (import cycles) and CodebaseDepPathTool (one dependency chain between two modules). Their parameter schemas are module constants. Other CodebaseTools queries are deliberately not wired here.

_LazyCodebaseTools defers construction of the backend to the first query, so building or registering the registry never touches Prolog. A construction failure is not cached or hidden: the real RuntimeError text appears in the failed ToolResult and the next call retries. _run_query is the shared body of execute, with _ok and _error building results. With no injected backend, _build_production_kb resolves the knowledge-base directory and AST snapshot from the environment, attaches to a fresh facts file or reindexes from the snapshot. register_codebase_query_tools registers the four tools and returns them; passing a kb bypasses auto-indexing.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

_LazyCodebaseTools

Line 131

Defers :class:CodebaseTools construction to the first .get() call.

CodebaseImpactTool

Bases: BaseTool · Line 218

Everything that transitively depends on a symbol (CodebaseTools.impact).

CodebaseModuleDepsTool

Bases: BaseTool · Line 281

Transitive internal module dependencies (CodebaseTools.module_deps).

CodebaseFindCyclesTool

Bases: BaseTool · Line 335

All import cycles among internal modules (CodebaseTools.cycles).

CodebaseDepPathTool

Bases: BaseTool · Line 386

One concrete dependency chain between modules (CodebaseTools.dep_path).

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/tools/codebase_query_tools.py

"""Registry-facing wrappers over CodebaseTools' Prolog codebase-intelligence
query surface -- SPEC-062/072 unified tool registry wiring.

:class:`~orchestrata_ultimate.tools.codebase_tools.CodebaseTools` (the
janus-backed Prolog codebase-intelligence layer) predates the unified tool
registry and was never registered into it. This module is the narrow "connect
A to B" wiring: four :class:`~orchestrata_ultimate.tools.base.BaseTool`
wrappers over CodebaseTools' impact / module-dependency / cycle / dependency-
path query predicates, following the same composition pattern as
``context_tools.py``'s ``GetCodeTool`` family (a private lazy backend wrapper,
one ``BaseTool`` subclass per query, a ``register_*`` function).

Fail-fast is preserved, not weakened: :class:`CodebaseTools`'s constructor
raises an actionable ``RuntimeError`` when the patched free-threaded
``janus-swi`` is unavailable or the process's GIL is enabled (see
``codebase_tools.py``'s module docstring) -- the same doctrine as the
kernel-required gate. :class:`_LazyCodebaseTools` defers that construction to
the first actual query (never at import or registry-registration time, so
merely building/registering the registry never touches Prolog), and -- unlike
``context_tools.py``'s ``_LazyRepo`` -- never caches or silences a
construction failure into a generic "unavailable" message: the real
``RuntimeError`` text is surfaced verbatim in the failed :class:`ToolResult`,
and every subsequent call retries construction rather than treating the
failure as permanent.

Query methods NOT wired here (``stats``/``find_symbol``/``symbol_context``/
``callers``/``callees``/``externals``/``unused_symbols``/``hotspots``/
``by_prop``/``signature_search``/``module_api``) are a known, deliberate
scope boundary of this wiring pass -- see the build brief.

The KB build/reindex lifecycle IS wired here: the shared :class:`CodebaseTools`
instance auto-indexes itself on first use (see
:func:`_build_production_kb`) -- a fresh ``.get()`` with no injected ``kb``
resolves the production KB directory + AST snapshot from the environment,
attaches to an already-fresh facts file, or reindexes from the snapshot when
none exists / it is stale. Passing ``kb=`` to
:func:`register_codebase_query_tools` (as every test in
``tests/tools/test_codebase_query_tools.py`` does) bypasses this entirely --
that instance is used as-is and auto-indexing never runs.
"""

from __future__ import annotations

import json
import os
import threading
from pathlib import Path
from typing import Any

from orchestrata_ultimate.core.paths import get_project_root
from orchestrata_ultimate.tools.base import (
    BaseTool,
    ToolCategory,
    ToolMetadata,
    ToolPriority,
    ToolResult,
)
from orchestrata_ultimate.tools.codebase_tools import CodebaseTools
from orchestrata_ultimate.tools.registry import ToolRegistry

# ---------------------------------------------------------------------------
# Production KB lifecycle -- env seams + auto-index/attach on first use.
# ---------------------------------------------------------------------------

# Relative to the project root (mirrors scripts/ast_daemon.py's
# ``.claude/run/`` precedent for a stable per-repo run directory).
_DEFAULT_KB_DIR = ".claude/run/codebase_kb"
# Relative to the project root -- the same snapshot `CLAUDE.md`'s AST-toolkit
# commands regenerate via ``serialize_codebase_ast.py serialize``.
_DEFAULT_AST_SNAPSHOT = "plans/ast-ultimate.json"


def _build_production_kb(environ: dict[str, str] | None) -> CodebaseTools:
    """Construct the shared production :class:`CodebaseTools` and index it.

    Resolves ``ORCHESTRATA_CODEBASE_KB_DIR``/``ORCHESTRATA_AST_SNAPSHOT`` via
    ``os.environ.get`` at call time (house rule: no bare env access), each
    defaulting to a stable path anchored at :func:`get_project_root`.
    Freshness is a plain mtime comparison: an existing ``facts.pl`` at least
    as new as the AST snapshot is attached to directly (skipping the ~2.17s
    subprocess reindex); otherwise (missing or stale) a real ``reindex()``
    runs once. A missing snapshot -- nothing to build the KB from -- raises
    an actionable ``RuntimeError`` (caught by :func:`_run_query`'s
    ``except RuntimeError``, so this never crashes an in-flight tool loop,
    matching the fail-fast-but-caught doctrine ``CodebaseTools.__init__``
    itself already uses for a missing janus/GIL-on environment).

    ``environ`` is the same test-injection seam as
    :func:`~orchestrata_ultimate.core.paths.get_project_root` -- pass an
    explicit dict to exercise this without touching real process env
    (monkeypatch is forbidden project-wide).
    """
    env = os.environ if environ is None else environ
    project_root = get_project_root(environ=environ)
    kb_dir = Path(
        env.get("ORCHESTRATA_CODEBASE_KB_DIR", str(project_root / _DEFAULT_KB_DIR))
    )
    snapshot_path = Path(
        env.get("ORCHESTRATA_AST_SNAPSHOT", str(project_root / _DEFAULT_AST_SNAPSHOT))
    )
    kb = CodebaseTools(workdir=str(kb_dir))
    facts_file = kb_dir / "facts.pl"

    if not snapshot_path.exists():
        raise RuntimeError(
            f"codebase KB snapshot not found at {snapshot_path} -- generate "
            "it with `uv run python scripts/serialize_codebase_ast.py "
            "serialize --source src/orchestrata_ultimate --name "
            f"orchestrata-ultimate --output {snapshot_path}`, or point "
            "ORCHESTRATA_AST_SNAPSHOT at an existing snapshot."
        )

    if (
        facts_file.exists()
        and facts_file.stat().st_mtime >= snapshot_path.stat().st_mtime
    ):
        kb.attach(str(facts_file))
        return kb

    result = kb.reindex(str(snapshot_path))
    if "error" in result:
        raise RuntimeError(f"codebase KB reindex failed: {result}")
    return kb


# ---------------------------------------------------------------------------
# Lazy, fail-fast-preserving backend handle
# ---------------------------------------------------------------------------


class _LazyCodebaseTools:
    """Defers :class:`CodebaseTools` construction to the first ``.get()`` call.

    Lets the four tools below be constructed with no arguments (so
    ``register_native_tools()`` — which builds every native tool eagerly on
    every registry build — never consults Prolog), while still sharing one
    instance across the four tools when built via
    :func:`register_codebase_query_tools`.

    Deliberately NOT the "degrade, don't crash" shape of ``context_tools.py``'s
    ``_LazyRepo``: a construction failure (missing janus / GIL enabled, or a
    production auto-index failure -- see :func:`_build_production_kb`) is
    never cached and never swallowed into a generic message -- ``.get()``
    re-raises the actionable ``RuntimeError`` on every attempt until
    construction actually succeeds (nothing is assigned to ``self._kb`` when
    the attempt raises, so the next call retries from scratch).
    """

    def __init__(
        self,
        kb: CodebaseTools | None = None,
        environ: dict[str, str] | None = None,
    ) -> None:
        self._kb = kb
        self._lock = threading.Lock()
        self._environ = environ

    def get(self) -> CodebaseTools:
        if self._kb is not None:
            return self._kb
        with self._lock:
            if self._kb is None:
                self._kb = _build_production_kb(self._environ)
            return self._kb


def _ok(tool_name: str, payload: dict[str, Any]) -> ToolResult:
    return ToolResult(tool_name, True, json.dumps(payload))


def _error(tool_name: str, message: str) -> ToolResult:
    return ToolResult(tool_name, False, "", message)


async def _run_query(
    tool_name: str, lazy_kb: _LazyCodebaseTools, query: Any
) -> ToolResult:
    """Shared execute() body: resolve the backend, run ``query``, wrap the result.

    ``query`` is a zero-arg callable (a closure over the already-validated
    arguments) so each tool's own argument parsing/validation stays local to
    its ``execute()`` while the backend-resolution + error-shaping is not
    duplicated four times.
    """
    try:
        kb = lazy_kb.get()
    except RuntimeError as exc:
        return _error(tool_name, str(exc))
    try:
        result = query(kb)
    except Exception as exc:  # noqa: BLE001 - never propagate to the caller
        return _error(tool_name, str(exc))
    if isinstance(result, dict) and "error" in result:
        return _error(tool_name, json.dumps(result))
    return _ok(tool_name, result)


# ---------------------------------------------------------------------------
# CodebaseImpactTool
# ---------------------------------------------------------------------------

_IMPACT_PARAMETERS: dict[str, Any] = {
    "type": "object",
    "properties": {
        "qname": {
            "type": "string",
            "description": "Fully-qualified symbol name to compute transitive impact for.",
        },
        "limit": {
            "type": "integer",
            "description": "Maximum affected symbols to return (default 25).",
        },
    },
    "required": ["qname"],
}


class CodebaseImpactTool(BaseTool):
    """Everything that transitively depends on a symbol (``CodebaseTools.impact``)."""

    def __init__(self, lazy_kb: _LazyCodebaseTools | None = None) -> None:
        self._metadata = ToolMetadata(
            name="codebase_impact",
            description=(
                "Find everything that transitively depends on a code symbol "
                "(class/function/method), via the Prolog codebase-intelligence "
                "database. Requires the KB to have been indexed (reindex())."
            ),
            category=ToolCategory.SEARCH,
            parameters=_IMPACT_PARAMETERS,
            priority=ToolPriority.P2,
            tags=("codebase", "prolog", "codeintel", "impact"),
        )
        self._lazy_kb = lazy_kb if lazy_kb is not None else _LazyCodebaseTools()

    @property
    def metadata(self) -> ToolMetadata:
        return self._metadata

    async def execute(self, arguments: dict[str, Any]) -> ToolResult:
        qname = arguments.get("qname")
        if not isinstance(qname, str) or not qname.strip():
            return _error(self.name, "qname is required")
        try:
            limit = int(arguments.get("limit", 25))
        except (TypeError, ValueError) as exc:
            return _error(self.name, f"invalid limit: {exc}")
        return await _run_query(
            self.name, self._lazy_kb, lambda kb: kb.impact(qname, limit=limit)
        )


# ---------------------------------------------------------------------------
# CodebaseModuleDepsTool
# ---------------------------------------------------------------------------

_MODULE_DEPS_PARAMETERS: dict[str, Any] = {
    "type": "object",
    "properties": {
        "module": {
            "type": "string",
            "description": "Internal module path (e.g. 'pkg.sub.mod') to query.",
        },
        "direction": {
            "type": "string",
            "enum": ["deps", "dependents"],
            "description": (
                "'deps' (default): modules this one transitively depends on. "
                "'dependents': modules that transitively depend on this one."
            ),
        },
        "limit": {
            "type": "integer",
            "description": "Maximum modules to return (default 50).",
        },
    },
    "required": ["module"],
}


class CodebaseModuleDepsTool(BaseTool):
    """Transitive internal module dependencies (``CodebaseTools.module_deps``)."""

    def __init__(self, lazy_kb: _LazyCodebaseTools | None = None) -> None:
        self._metadata = ToolMetadata(
            name="codebase_module_deps",
            description=(
                "List a module's transitive internal dependencies (or, with "
                "direction='dependents', everything that depends on it), via "
                "the Prolog codebase-intelligence database."
            ),
            category=ToolCategory.SEARCH,
            parameters=_MODULE_DEPS_PARAMETERS,
            priority=ToolPriority.P2,
            tags=("codebase", "prolog", "codeintel", "modules"),
        )
        self._lazy_kb = lazy_kb if lazy_kb is not None else _LazyCodebaseTools()

    @property
    def metadata(self) -> ToolMetadata:
        return self._metadata

    async def execute(self, arguments: dict[str, Any]) -> ToolResult:
        module = arguments.get("module")
        if not isinstance(module, str) or not module.strip():
            return _error(self.name, "module is required")
        direction = arguments.get("direction", "deps")
        try:
            limit = int(arguments.get("limit", 50))
        except (TypeError, ValueError) as exc:
            return _error(self.name, f"invalid limit: {exc}")
        return await _run_query(
            self.name,
            self._lazy_kb,
            lambda kb: kb.module_deps(module, direction=direction, limit=limit),
        )


# ---------------------------------------------------------------------------
# CodebaseFindCyclesTool
# ---------------------------------------------------------------------------

_CYCLES_PARAMETERS: dict[str, Any] = {
    "type": "object",
    "properties": {
        "limit": {
            "type": "integer",
            "description": "Maximum import cycles to return (default 20).",
        },
    },
    "required": [],
}


class CodebaseFindCyclesTool(BaseTool):
    """All import cycles among internal modules (``CodebaseTools.cycles``)."""

    def __init__(self, lazy_kb: _LazyCodebaseTools | None = None) -> None:
        self._metadata = ToolMetadata(
            name="codebase_find_cycles",
            description=(
                "Find import cycles among internal modules, via the Prolog "
                "codebase-intelligence database."
            ),
            category=ToolCategory.SEARCH,
            parameters=_CYCLES_PARAMETERS,
            priority=ToolPriority.P2,
            tags=("codebase", "prolog", "codeintel", "cycles"),
        )
        self._lazy_kb = lazy_kb if lazy_kb is not None else _LazyCodebaseTools()

    @property
    def metadata(self) -> ToolMetadata:
        return self._metadata

    async def execute(self, arguments: dict[str, Any]) -> ToolResult:
        try:
            limit = int(arguments.get("limit", 20))
        except (TypeError, ValueError) as exc:
            return _error(self.name, f"invalid limit: {exc}")
        return await _run_query(
            self.name, self._lazy_kb, lambda kb: kb.cycles(limit=limit)
        )


# ---------------------------------------------------------------------------
# CodebaseDepPathTool
# ---------------------------------------------------------------------------

_DEP_PATH_PARAMETERS: dict[str, Any] = {
    "type": "object",
    "properties": {
        "from_module": {
            "type": "string",
            "description": "Source internal module path.",
        },
        "to_module": {
            "type": "string",
            "description": "Target internal module path.",
        },
    },
    "required": ["from_module", "to_module"],
}


class CodebaseDepPathTool(BaseTool):
    """One concrete dependency chain between modules (``CodebaseTools.dep_path``)."""

    def __init__(self, lazy_kb: _LazyCodebaseTools | None = None) -> None:
        self._metadata = ToolMetadata(
            name="codebase_dep_path",
            description=(
                "Find one concrete import-dependency chain from one internal "
                "module to another, via the Prolog codebase-intelligence "
                "database."
            ),
            category=ToolCategory.SEARCH,
            parameters=_DEP_PATH_PARAMETERS,
            priority=ToolPriority.P2,
            tags=("codebase", "prolog", "codeintel", "dependency-path"),
        )
        self._lazy_kb = lazy_kb if lazy_kb is not None else _LazyCodebaseTools()

    @property
    def metadata(self) -> ToolMetadata:
        return self._metadata

    async def execute(self, arguments: dict[str, Any]) -> ToolResult:
        from_module = arguments.get("from_module")
        to_module = arguments.get("to_module")
        if not isinstance(from_module, str) or not from_module.strip():
            return _error(self.name, "from_module is required")
        if not isinstance(to_module, str) or not to_module.strip():
            return _error(self.name, "to_module is required")
        return await _run_query(
            self.name,
            self._lazy_kb,
            lambda kb: kb.dep_path(from_module, to_module),
        )


# ---------------------------------------------------------------------------
# Wiring
# ---------------------------------------------------------------------------


def register_codebase_query_tools(
    registry: ToolRegistry,
    kb: CodebaseTools | None = None,
    environ: dict[str, str] | None = None,
) -> tuple[
    CodebaseImpactTool,
    CodebaseModuleDepsTool,
    CodebaseFindCyclesTool,
    CodebaseDepPathTool,
]:
    """Register the four Prolog codebase-intelligence query tools.

    All four share ONE lazily-constructed :class:`CodebaseTools` instance
    (unless ``kb`` is given, in which case that instance is used directly and
    construction never happens here) -- registering never constructs
    ``CodebaseTools`` and never consults Prolog. ``environ`` is the test-
    injection seam for :func:`_build_production_kb`'s env resolution
    (``ORCHESTRATA_CODEBASE_KB_DIR``/``ORCHESTRATA_AST_SNAPSHOT``); it is
    ignored whenever ``kb`` is given.
    """
    lazy_kb = _LazyCodebaseTools(kb, environ=environ)
    impact_tool = CodebaseImpactTool(lazy_kb)
    module_deps_tool = CodebaseModuleDepsTool(lazy_kb)
    find_cycles_tool = CodebaseFindCyclesTool(lazy_kb)
    dep_path_tool = CodebaseDepPathTool(lazy_kb)
    registry.register(impact_tool, replace=True)
    registry.register(module_deps_tool, replace=True)
    registry.register(find_cycles_tool, replace=True)
    registry.register(dep_path_tool, replace=True)
    return impact_tool, module_deps_tool, find_cycles_tool, dep_path_tool