codex-app-server/cloud/tree.py

"""任务 worktree 的结构化目录树（server / MCP / worker 首轮注入共用）。

只做一件事：把一个 worktree（或其子目录）列成结构化树，
不读文件内容、不跟符号链接、不出 worktree。

安全规则与 /files 端点同款：
- 相对路径 only：绝对路径、.. 逃逸、符号链接逃逸一律拒绝；
  目录内的符号链接列为 type=symlink、size=0、不跟进、不展开。
- .git 永远不列（即使调用方显式传 ignore 也强制加上）；
  默认还忽略 node_modules、__pycache__、.venv、venv 等
  构建/依赖噪音目录。.gitignore 不解析——解析会随仓库
  变化而让同一调用返回不同结构，显式忽略集更可预测。
- 条数封顶：遍历按“目录在前、同类按名排序”的确定顺序，
  到达 max_entries 即停并置 truncated=True，绝不静默少列。

返回形状：
  {"path": <相对路径，根为 "">, "truncated": bool,
   "entries": [{"name", "path", "type": "dir"|"file"|"symlink",
                "size": int（文件才有意义，目录/链接为 0）,
                "children": [...]（仅 dir 有；
                被 max_depth 截断的目录 children=[] 且
                children_truncated=True）}]}
max_depth 从请求根算起：1 只列根的直接条目，2 再列一层
子目录，以此类推；默认 3。render_tree_text 把同一结构
渲染成首轮注入用的纯文本树（只含名、目录斜杠与链接 @，
不含大小，省 token）。
"""

import os

__all__ = [
    "DEFAULT_IGNORE",
    "DEFAULT_MAX_DEPTH",
    "DEFAULT_MAX_ENTRIES",
    "INJECTION_MAX_DEPTH",
    "INJECTION_MAX_ENTRIES",
    "build_tree",
    "render_tree_text",
    "resolve_tree_path",
]

DEFAULT_IGNORE = frozenset({
    ".git",
    "node_modules",
    "__pycache__",
    ".venv",
    "venv",
    ".tox",
    ".eggs",
})
DEFAULT_MAX_DEPTH = 3
DEFAULT_MAX_ENTRIES = 200
# 首轮注入比 API 更紧：注入是每个任务都付的 prompt 成本，
# 2 层 + 120 条足够定位， deeper 靠 shell 或 get_task_tree 再查。
INJECTION_MAX_DEPTH = 2
INJECTION_MAX_ENTRIES = 120


def resolve_tree_path(root, rel):
    """worktree 内相对路径的安全解析：逃逸返回 None。

    空串返回 root 本身；绝对路径、.. 逃逸、符号链接逃逸
    （realpath 后越界）一律 None。目录内指回内部的链接
    放行到其真实位置——调用方仍只拿到 worktree 内的内容，
    但 build_tree 遇到链接本身时不跟进（见下），所以这条
    放行只影响“显式 path 指向内部链接目录”时的根解析。
    """
    if not rel:
        return root
    if os.path.isabs(rel):
        return None
    base = os.path.realpath(root)
    target = os.path.realpath(os.path.join(base, rel))
    try:
        if os.path.commonpath([base, target]) != base:
            return None
    except ValueError:
        return None
    return target


def build_tree(root, path="", max_depth=DEFAULT_MAX_DEPTH,
               max_entries=DEFAULT_MAX_ENTRIES, ignore=None):
    """列 root 下 path 子目录的结构化树。

    path 为空时列 root 本身。path 逃逸时抛 ValueError；
    目标不存在时抛 FileNotFoundError；目标不是目录时抛
    NotADirectoryError——调用方按各自的 HTTP 语义映射
    （server：逃逸 400、不存在/非目录 404）。
    """
    target = resolve_tree_path(root, path)
    if target is None:
        raise ValueError("path escapes task worktree")
    if not os.path.exists(target):
        raise FileNotFoundError("no such path: %s" % path)
    if not os.path.isdir(target):
        raise NotADirectoryError("path is not a directory: %s" % path)
    effective_ignore = set(ignore) if ignore is not None \
        else set(DEFAULT_IGNORE)
    effective_ignore.add(".git")
    state = {"count": 0, "truncated": False}
    entries = _walk(target, root, path, depth=1,
                    max_depth=max(1, int(max_depth)),
                    max_entries=max(1, int(max_entries)),
                    ignore=effective_ignore, state=state)
    return {"path": path, "truncated": state["truncated"],
            "entries": entries}


def _walk(dir_path, root, rel_prefix, depth, max_depth, max_entries,
          ignore, state):
    if depth > max_depth:
        return []
    try:
        names = os.listdir(dir_path)
    except OSError:
        return []
    dirs = []
    others = []
    for name in names:
        if name in ignore:
            continue
        full = os.path.join(dir_path, name)
        if os.path.islink(full):
            others.append(name)
        elif os.path.isdir(full):
            dirs.append(name)
        else:
            others.append(name)
    # 目录在前、同类按名排序：遍历顺序确定，截断结果可复现。
    ordered = sorted(dirs, key=str.lower) + sorted(others, key=str.lower)
    entries = []
    for name in ordered:
        if state["count"] >= max_entries:
            state["truncated"] = True
            break
        full = os.path.join(dir_path, name)
        rel = (rel_prefix + "/" + name).lstrip("/") \
            if rel_prefix else name
        # Windows 风格分隔符不进结构：统一用 /。
        rel = rel.replace(os.sep, "/")
        state["count"] += 1
        if os.path.islink(full):
            entries.append({"name": name, "path": rel,
                            "type": "symlink", "size": 0})
        elif os.path.isdir(full):
            if depth >= max_depth:
                entries.append({"name": name, "path": rel,
                                "type": "dir", "size": 0,
                                "children": [],
                                "children_truncated": True})
            else:
                children = _walk(full, root, rel, depth + 1,
                                 max_depth, max_entries, ignore, state)
                entries.append({"name": name, "path": rel,
                                "type": "dir", "size": 0,
                                "children": children,
                                # 条数截断发生在子层时同样要标出来，
                                # 否则半截 children 会被当成完整目录。
                                "children_truncated": state["truncated"]})
        else:
            try:
                size = os.path.getsize(full)
            except OSError:
                size = 0
            entries.append({"name": name, "path": rel,
                            "type": "file", "size": size})
        if state["truncated"]:
            break
    return entries


def render_tree_text(tree):
    """把 build_tree 的结果渲染成纯文本树（首轮注入用）。

    只含名与目录斜杠，不含大小；截断时末尾明写“已截断”。
    """
    lines = ["."]
    _render_entries(tree.get("entries") or [], "", lines)
    if tree.get("truncated"):
        lines.append("（以上已达条数上限，未列全；深层结构用 shell "
                     "或 get_task_tree 再查。）")
    return "\n".join(lines)


def _render_entries(entries, prefix, lines):
    for index, entry in enumerate(entries):
        last = index == len(entries) - 1
        branch = "└── " if last else "├── "
        suffix = "/" if entry.get("type") == "dir" else (
            "@" if entry.get("type") == "symlink" else "")
        lines.append(prefix + branch + entry["name"] + suffix)
        children = entry.get("children") or []
        if entry.get("type") == "dir" and children:
            extension = "    " if last else "│   "
            _render_entries(children, prefix + extension, lines)
        elif entry.get("type") == "dir" \
                and entry.get("children_truncated"):
            extension = "    " if last else "│   "
            lines.append(prefix + extension + "…")