LLeMbas/src/lembas/services/agent/index.py

"""What is in a project directory, for the picker and for the model.

Two things want this list. The `@` picker needs something to filter, and a
model working in a directory should know roughly what is in it rather than
spending its first two rounds finding out. Both want the same walk, so it
happens once and is cached.

**Three ways of getting it, in order.** `git ls-files` first, because most
project directories are repositories and it applies `.gitignore` for free --
without which the answer for a Node project is forty thousand paths under
`node_modules`. Then `find`, with the usual noise pruned by hand. Then a
recursive SFTP walk, which always works and costs a round trip per directory.

**Two commands run here, and neither goes through `agent/policy.py`.** That is
deliberate and it is the same argument the terminal panel and the directory
browser rest on: this is LLeMbas listing a directory on somebody's behalf, not
a model choosing to run something. Both are read-only, both are built here
rather than assembled from anything a model said, and the project directory is
configuration rather than input. It is still an exception to Manual mode's
"everything is shown to you before it happens", and it is written down in
CLAUDE.md next to the others.

**Nothing here is trusted.** Filenames come off somebody else's machine and end
up inside a system prompt, so they are stripped of control characters, capped
in length, capped in number, and never interpreted.
"""

from __future__ import annotations

import asyncio
import logging
import re
import time
from dataclasses import dataclass, field

from lembas.services.agent.base import ExecError, ExecRequest, Executor

log = logging.getLogger(__name__)

# How many paths are kept. Past this the index says it was truncated, which the
# rendering repeats to the model -- "there is nothing else here" and "I stopped
# looking" are different answers and it must not give the first for the second.
MAX_ENTRIES = 20_000
# One path. Longer than any real one and shorter than an attack.
MAX_PATH = 400
# How long a walk may take before it is abandoned. The index is a convenience;
# a chat must never sit waiting for one.
BUILD_TIMEOUT = 20.0
# Output budget for the listing commands. Twenty thousand paths at forty
# characters is 800KB, so this has room and still refuses a runaway.
MAX_OUTPUT = 2 * 1024 * 1024

# How long a built index is reused, and how many are kept at once. A project
# directory changes under you -- the model writes files into it -- so this is
# short. `refresh` exists for when short is not short enough.
TTL = 300.0
MAX_CACHED = 64

# How deep the SFTP fallback goes, and how many directories it will open. It is
# a round trip per directory, so an unbounded walk of somebody's home directory
# would take minutes and achieve nothing.
SFTP_MAX_DEPTH = 6
SFTP_MAX_DIRS = 400

# Pruned from the `find` and SFTP paths. Not applied to `git ls-files`, which
# has already applied the repository's own rules and where a checked-in
# `vendor/` is checked in on purpose -- this project's own hash-pinned browser
# libraries live in one.
IGNORED = (
    ".git",
    ".hg",
    ".svn",
    "node_modules",
    "__pycache__",
    ".venv",
    "venv",
    ".mypy_cache",
    ".pytest_cache",
    ".ruff_cache",
    ".tox",
    ".next",
    ".nuxt",
    ".gradle",
    ".terraform",
    "target",
    "dist",
    "build",
    ".DS_Store",
)

# Control characters, including the escape that would let a filename repaint
# the transcript it is quoted in.
_CONTROL = re.compile(r"[\x00-\x1f\x7f-\x9f]")


@dataclass(frozen=True)
class ProjectIndex:
    """A snapshot of what was in a directory, and how it was found out."""

    paths: tuple[str, ...] = ()
    total: int = 0
    truncated: bool = False
    source: str = ""
    built_at: float = field(default=0.0)

    @property
    def ok(self) -> bool:
        return bool(self.paths)


# --- Building ----------------------------------------------------------------
def _clean(raw: str) -> str:
    """One path, made safe to put in a prompt and in an attribute."""
    path = _CONTROL.sub("", raw.strip()).lstrip("./")
    return path[:MAX_PATH]


def _collect(output: str) -> tuple[tuple[str, ...], int, bool]:
    seen: set[str] = set()
    paths: list[str] = []
    total = 0
    for line in output.splitlines():
        path = _clean(line)
        if not path or path in seen:
            continue
        total += 1
        if len(paths) < MAX_ENTRIES:
            seen.add(path)
            paths.append(path)
    paths.sort()
    return tuple(paths), total, total > len(paths)


async def _from_git(executor: Executor, project_dir: str) -> ProjectIndex | None:
    """Tracked and untracked files, minus whatever `.gitignore` excludes.

    `--exclude-standard` is what makes this worth trying first: the repository
    already carries somebody's considered list of what is not part of the
    project, and reproducing it by hand is how an index ends up ninety percent
    build output.
    """
    result = await executor.run(
        ExecRequest(
            command="git ls-files -c -o --exclude-standard 2>/dev/null",
            cwd=project_dir,
            timeout=BUILD_TIMEOUT,
            max_bytes=MAX_OUTPUT,
        )
    )
    if not result.ok or not result.output.strip():
        return None
    paths, total, truncated = _collect(result.output)
    if not paths:
        return None
    return ProjectIndex(
        paths=paths, total=total, truncated=truncated or result.truncated, source="git"
    )


def _find_command() -> str:
    prunes = " -o ".join(f"-name {name!r}" for name in IGNORED)
    # -print rather than -print0: the output is read as text either way, and a
    # filename containing a newline splits into two entries that resolve to
    # nothing rather than into anything dangerous.
    return f"find . \\( {prunes} \\) -prune -o -print 2>/dev/null"


async def _from_find(executor: Executor, project_dir: str) -> ProjectIndex | None:
    result = await executor.run(
        ExecRequest(
            command=_find_command(),
            cwd=project_dir,
            timeout=BUILD_TIMEOUT,
            max_bytes=MAX_OUTPUT,
        )
    )
    if not result.output.strip():
        return None
    paths, total, truncated = _collect(result.output)
    if not paths:
        return None
    return ProjectIndex(
        paths=paths, total=total, truncated=truncated or result.truncated, source="find"
    )


async def _from_sftp(executor: Executor, project_dir: str) -> ProjectIndex:
    """The one that always works, and the one that is slow.

    Bounded twice over -- by depth and by how many directories it will open --
    because this is a network round trip per directory and an unbounded walk of
    a home directory would take minutes to produce something unusable.
    """
    found: list[str] = []
    opened = 0
    queue: list[tuple[str, int]] = [("", 0)]

    while queue and opened < SFTP_MAX_DIRS and len(found) < MAX_ENTRIES:
        where, depth = queue.pop(0)
        opened += 1
        try:
            entries = await executor.scan_dir(where or project_dir)
        except ExecError:
            continue
        for entry in entries:
            if entry.name in IGNORED:
                continue
            path = f"{where}/{entry.name}" if where else entry.name
            found.append(path + "/" if entry.is_dir else path)
            if entry.is_dir and depth + 1 < SFTP_MAX_DEPTH:
                queue.append((path, depth + 1))

    paths, total, truncated = _collect("\n".join(found))
    return ProjectIndex(
        paths=paths,
        total=total,
        truncated=truncated or bool(queue),
        source="sftp",
    )


async def build(executor: Executor, project_dir: str) -> ProjectIndex:
    """Walk the directory, by whichever means works first."""
    started = time.monotonic()
    try:
        found = None
        for attempt in (_from_git, _from_find):
            try:
                found = await attempt(executor, project_dir)
            except ExecError as exc:
                # A rung that cannot run at all is a rung that did not answer,
                # not the end of the ladder. A host that refuses exec entirely
                # -- an SFTP-only account, a forced command -- is the exact case
                # the SFTP rung below exists for, and letting this out skipped
                # straight past it to an empty listing.
                log.debug("indexing %s: %s did not run: %s", project_dir, attempt.__name__,
                          exc.message)
                found = None
            if found is not None:
                break
        if found is None:
            found = await _from_sftp(executor, project_dir)
    except ExecError as exc:
        log.info("could not index %s: %s", project_dir, exc.message)
        return ProjectIndex(built_at=time.monotonic())

    log.debug(
        "indexed %s: %d paths by %s in %dms",
        project_dir,
        len(found.paths),
        found.source,
        int((time.monotonic() - started) * 1000),
    )
    return ProjectIndex(
        paths=found.paths,
        total=found.total,
        truncated=found.truncated,
        source=found.source,
        built_at=time.monotonic(),
    )


# --- The cache ---------------------------------------------------------------
# Keyed on the connection and the directory, not the chat: two chats on the same
# box in the same tree are looking at the same files, and indexing it twice
# would double the cost to prove it.
_CACHE: dict[tuple[str, str], ProjectIndex] = {}
_BUILDING: dict[tuple[str, str], asyncio.Task] = {}


def cached(profile_id: str, project_dir: str) -> ProjectIndex | None:
    """What is already known, or None. Never does any work.

    `harness.context_variables` is synchronous and sits on the request path, so
    it may only ever call this -- an SFTP round trip from there would block a
    request while somebody's box thought about it.
    """
    found = _CACHE.get((profile_id, project_dir))
    if found is None:
        return None
    if time.monotonic() - found.built_at > TTL:
        _CACHE.pop((profile_id, project_dir), None)
        return None
    return found


async def ensure(
    executor: Executor, profile_id: str, project_dir: str, *, refresh: bool = False
) -> ProjectIndex:
    """The index, building it if there is not a fresh one already.

    Concurrent callers share one build. A reply and the `@` picker asking at
    the same moment is the ordinary case, not a rare one, and two walks of the
    same tree would be two of everything for one answer.
    """
    key = (profile_id, project_dir)
    if refresh:
        _CACHE.pop(key, None)
    elif (found := cached(profile_id, project_dir)) is not None:
        return found

    if (running := _BUILDING.get(key)) is not None:
        return await asyncio.shield(running)

    task = asyncio.create_task(build(executor, project_dir))
    _BUILDING[key] = task
    try:
        found = await task
    finally:
        _BUILDING.pop(key, None)

    _CACHE[key] = found
    while len(_CACHE) > MAX_CACHED:
        _CACHE.pop(next(iter(_CACHE)))
    return found


# --- Rendering ---------------------------------------------------------------
# A tree that lists a thousand files is worse than no tree: it costs the window
# on every request forever and buries the four names that mattered. So the
# rendering has a character budget and elides what will not fit, saying how much
# it elided -- a directory shown as `src/vendor/ (412 files)` is a model being
# told where to look, which is the useful half of listing it.
INDENT = "  "
# Below this a directory is never collapsed. Elision costs a line either way, so
# collapsing three files into "(3 files)" saves nothing and loses everything.
ALWAYS_SHOW = 4


def _tree(paths: tuple[str, ...]) -> dict:
    root: dict = {}
    for path in paths:
        node = root
        parts = [part for part in path.rstrip("/").split("/") if part]
        for part in parts[:-1]:
            node = node.setdefault(part, {})
            if not isinstance(node, dict):  # a file and a directory share a name
                break
        else:
            if parts:
                leaf = parts[-1]
                if path.endswith("/"):
                    node.setdefault(leaf, {})
                else:
                    node.setdefault(leaf, None)
    return root


def _files_under(node: dict) -> int:
    total = 0
    for child in node.values():
        total += _files_under(child) if isinstance(child, dict) else 1
    return total


def _candidates(node: dict, prefix: str, depth: int, out: list) -> None:
    """Every directory, with what collapsing it would save."""
    for name, child in node.items():
        if not isinstance(child, dict):
            continue
        path = f"{prefix}{name}/"
        count = _files_under(child)
        full = _cost(child, depth + 1)
        collapsed = len(f" ({count} files)")
        if count > ALWAYS_SHOW and full > collapsed:
            out.append((depth, count, path, full - collapsed))
        _candidates(child, path, depth + 1, out)


def _cost(node: dict, depth: int) -> int:
    """Roughly how many characters rendering this subtree in full would take."""
    total = 0
    for name, child in node.items():
        total += len(INDENT) * (depth + 1) + len(name) + 2
        if isinstance(child, dict):
            total += _cost(child, depth + 1)
    return total


def _plan(root: dict, budget: int) -> set[str]:
    """Which directories to show as a count, so the rest fits.

    Deepest and largest first. Collapsing by saving alone would take `src/`
    before `src/web/static/vendor/` -- it is bigger, because it *contains* it --
    and lose every name worth having to save one directory of hash-pinned
    third-party files. Depth is the proxy for "further from what somebody was
    looking for", and it is a good one.
    """
    if _cost(root, 0) <= budget:
        return set()

    candidates: list[tuple[int, int, str, int]] = []
    _candidates(root, "", 0, candidates)
    candidates.sort(key=lambda item: (-item[0], -item[1]))

    chosen: dict[str, int] = {}
    saved = 0
    total = _cost(root, 0)
    for _depth, _count, path, saving in candidates:
        if total - saved <= budget:
            break
        # A directory inside one already collapsed is not rendered at all, so
        # collapsing it saves nothing.
        if any(path.startswith(done) for done in chosen):
            continue
        # And a directory *containing* one already collapsed subsumes it. Its
        # own saving is measured against the full subtree, so the descendant's
        # has to come back off or the two are counted twice -- which stopped
        # the loop early believing it had made room it had not.
        for inside in [done for done in chosen if done.startswith(path)]:
            saved -= chosen.pop(inside)
        chosen[path] = saving
        saved += saving
    return set(chosen)


def _lines(
    node: dict, prefix: str, depth: int, collapsed: set[str], budget: list[int]
) -> list[str]:
    out: list[str] = []
    # Files before directories at each level: the shallow names are the ones
    # somebody would recognise, and if the budget runs out mid-tree they are
    # the ones worth having spent it on.
    files = sorted(name for name, child in node.items() if not isinstance(child, dict))
    folders = sorted(name for name, child in node.items() if isinstance(child, dict))

    for position, name in enumerate(files):
        line = f"{INDENT * depth}{name}"
        if budget[0] < len(line) + 1:
            out.append(f"{INDENT * depth}… {len(files) - position} more files")
            budget[0] = 0
            return out
        budget[0] -= len(line) + 1
        out.append(line)

    for name in folders:
        child = node[name]
        path = f"{prefix}{name}/"
        header = f"{INDENT * depth}{name}/"
        if path in collapsed:
            line = f"{header} ({_files_under(child)} files)"
            budget[0] -= len(line) + 1
            out.append(line)
            continue
        if budget[0] < len(header) + 1:
            return out
        budget[0] -= len(header) + 1
        out.append(header)
        out.extend(_lines(child, path, depth + 1, collapsed, budget))
    return out


def render(index: ProjectIndex, budget: int) -> str:
    """The listing as the model sees it, inside `budget` characters.

    Returns "" when there is nothing to say, so the fragment carrying it can
    vanish entirely rather than appear as an empty heading -- which is what
    `Fragment.requires` is for.
    """
    if not index.ok or budget <= 0:
        return ""

    root = _tree(index.paths)
    collapsed = _plan(root, budget)
    # The plan has already made it fit, so this is a backstop rather than the
    # mechanism -- with enough slack that an estimate a little off does not
    # truncate a listing that was fine. What it is really for is the one shape
    # collapsing cannot help with: five thousand files directly in the root,
    # where there is no directory to fold them into.
    remaining = [int(budget * 1.5) + 200]
    lines = _lines(root, "", 0, collapsed, remaining)
    if not lines:
        return ""

    note = ""
    if index.truncated:
        note = (
            f"\n\nThere are more than {len(index.paths)} entries here; this is the "
            "first of them, so treat it as a sample rather than the whole tree."
        )
    elif collapsed:
        note = (
            "\n\nDirectories shown with a count were left unopened to save room. "
            "Use `file_list` to look inside one."
        )
    return "\n".join(lines) + note


def forget(profile_id: str) -> int:
    """Drop everything indexed through one connection.

    Called when a profile is deleted, disabled or has its host key forgotten --
    the same moments that close its terminals. Keeping a listing of a machine
    somebody has just revoked would be a small leak of exactly the kind the
    rest of this module is careful about.
    """
    doomed = [key for key in _CACHE if key[0] == profile_id]
    for key in doomed:
        _CACHE.pop(key, None)
    return len(doomed)


def forget_dir(profile_id: str, project_dir: str) -> None:
    """Drop one tree's listing, because something just changed it.

    The TTL exists for drift nobody can see coming. A write through `file_write`
    is not that: it is this process changing the tree it has just described, and
    leaving five minutes of a listing that is known to be wrong is worse than
    having none -- a model reading it concludes the file it created is missing.
    """
    _CACHE.pop((profile_id, project_dir), None)


def clear() -> None:
    _CACHE.clear()


__all__ = [
    "MAX_ENTRIES",
    "ProjectIndex",
    "build",
    "cached",
    "clear",
    "ensure",
    "forget",
    "forget_dir",
    "render",
]