zai-cli 0.1.0__py3-none-any.whl

zai/core/repomap.py

@@ -0,0 +1,381 @@
+"""Fast, bounded, gitignore-aware repository mapping."""
+from __future__ import annotations
+
+import ast
+import hashlib
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+from .context import estimate_text_tokens
+from .storage import atomic_write_json, read_json
+
+SKIP_DIRS = {
+    ".git", ".zai", "__pycache__", "node_modules", ".venv", "venv",
+    "dist", "build", ".next", ".nuxt", "coverage", ".pytest_cache",
+    ".mypy_cache", ".ruff_cache", "target", "vendor",
+}
+SKIP_EXTS = {
+    ".pyc", ".pyo", ".lock", ".sum", ".png", ".jpg", ".jpeg", ".gif",
+    ".ico", ".svg", ".woff", ".woff2", ".ttf", ".eot", ".mp4", ".zip",
+    ".tar", ".gz", ".7z", ".pdf", ".exe", ".dll", ".so", ".dylib",
+}
+CODE_EXTS = {
+    ".py", ".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx", ".go", ".rs",
+    ".java", ".cpp", ".cc", ".c", ".h", ".hpp", ".cs", ".rb", ".php",
+    ".swift", ".kt", ".kts", ".vue", ".svelte", ".html", ".css", ".scss",
+    ".sql", ".sh", ".ps1", ".toml", ".yaml", ".yml", ".json", ".md",
+}
+HIGH_PRIORITY_NAMES = {
+    "readme.md", "pyproject.toml", "package.json", "cargo.toml", "go.mod",
+    "requirements.txt", "dockerfile", "compose.yml", "compose.yaml",
+    "makefile", "justfile",
+}
+SECRET_NAMES = {
+    ".env", ".env.local", ".env.production", "credentials.json",
+    "secrets.json", "id_rsa", "id_ed25519",
+}
+SECRET_SUFFIXES = {".pem", ".key", ".p12", ".pfx"}
+
+DEFAULT_MAX_FILES = 500
+DEFAULT_MAX_SCAN_FILES = 10_000
+DEFAULT_MAX_FILE_BYTES = 512 * 1024
+DEFAULT_MAX_REPO_BYTES = 100 * 1024 * 1024
+DEFAULT_PROMPT_TOKENS = 12_000
+CACHE_VERSION = 1
+
+
+@dataclass
+class RepoMapStats:
+    scanned: int = 0
+    indexed: int = 0
+    ignored: int = 0
+    unsupported: int = 0
+    oversized: int = 0
+    binary: int = 0
+    symlinks: int = 0
+    errors: int = 0
+    cache_hits: int = 0
+    cache_misses: int = 0
+    truncated: bool = False
+    bytes_considered: int = 0
+
+
+@dataclass
+class RepoMapResult:
+    text: str
+    stats: RepoMapStats
+    files: list[str]
+
+
+def _extract_python_symbols(path: str) -> list[str]:
+    try:
+        tree = ast.parse(Path(path).read_text(encoding="utf-8", errors="ignore"))
+        symbols = []
+        for node in tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                symbols.append(f"def {node.name}()")
+            elif isinstance(node, ast.ClassDef):
+                methods = [
+                    f".{item.name}()"
+                    for item in node.body
+                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
+                ]
+                suffix = f": {', '.join(methods[:3])}" if methods else ""
+                symbols.append(f"class {node.name}{suffix}")
+        return symbols[:12]
+    except (OSError, SyntaxError, UnicodeError):
+        return []
+
+
+def _extract_regex_symbols(content: str, suffix: str) -> list[str]:
+    patterns = {
+        ".go": r"^\s*(?:func|type)\s+([A-Za-z_]\w*)",
+        ".rs": r"^\s*(?:pub\s+)?(?:fn|struct|enum|trait|impl)\s+([A-Za-z_]\w*)",
+        ".java": r"^\s*(?:public|private|protected)?\s*(?:class|interface|enum)\s+([A-Za-z_]\w*)",
+        ".cs": r"^\s*(?:public|private|internal|protected)?\s*(?:class|interface|enum|record)\s+([A-Za-z_]\w*)",
+        ".rb": r"^\s*(?:def|class|module)\s+([A-Za-z_]\w*[!?=]?)",
+        ".php": r"^\s*(?:public\s+|private\s+|protected\s+)?(?:function|class|interface|trait)\s+([A-Za-z_]\w*)",
+        ".swift": r"^\s*(?:func|class|struct|enum|protocol)\s+([A-Za-z_]\w*)",
+        ".kt": r"^\s*(?:fun|class|object|interface|data\s+class)\s+([A-Za-z_]\w*)",
+        ".kts": r"^\s*(?:fun|class|object|interface|data\s+class)\s+([A-Za-z_]\w*)",
+    }
+    pattern = patterns.get(suffix)
+    if not pattern:
+        return []
+    return [
+        match.group(0).strip()[:100]
+        for match in re.finditer(pattern, content, re.MULTILINE)
+    ][:12]
+
+
+def _extract_js_symbols(content: str) -> list[str]:
+    pattern = re.compile(
+        r"^\s*(?:export\s+(?:default\s+)?)?"
+        r"(?:async\s+)?(?:function|class|const|let)\s+([A-Za-z_$][\w$]*)",
+        re.MULTILINE,
+    )
+    return [match.group(0).strip()[:100] for match in pattern.finditer(content)][:12]
+
+
+def _extract_symbols(path: Path, content: str | None = None) -> list[str]:
+    suffix = path.suffix.lower()
+    if suffix == ".py":
+        return _extract_python_symbols(str(path))
+    if content is None:
+        try:
+            content = path.read_text(encoding="utf-8", errors="ignore")
+        except OSError:
+            return []
+    if suffix in {".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx", ".vue", ".svelte"}:
+        return _extract_js_symbols(content)
+    return _extract_regex_symbols(content, suffix)
+
+
+def _is_secret(path: Path) -> bool:
+    name = path.name.lower()
+    return (
+        name in SECRET_NAMES
+        or name.startswith(".env.")
+        or path.suffix.lower() in SECRET_SUFFIXES
+        or any(part.lower() in {".ssh", ".aws", ".gnupg"} for part in path.parts)
+    )
+
+
+def _looks_binary(path: Path) -> bool:
+    try:
+        with path.open("rb") as handle:
+            chunk = handle.read(4096)
+        return b"\0" in chunk
+    except OSError:
+        return True
+
+
+def _load_gitignore(root: Path, extra_patterns: list[str] | None = None):
+    try:
+        import pathspec
+    except ImportError:
+        return None
+    patterns = []
+    for filename in (".gitignore", ".ignore"):
+        path = root / filename
+        if path.is_file():
+            try:
+                patterns.extend(path.read_text(encoding="utf-8", errors="ignore").splitlines())
+            except OSError:
+                pass
+    patterns.extend(extra_patterns or [])
+    return pathspec.GitIgnoreSpec.from_lines(patterns)
+
+
+def _ignored(spec, relative: str, is_dir: bool = False) -> bool:
+    if spec is None:
+        return False
+    candidate = relative.replace("\\", "/") + ("/" if is_dir else "")
+    return spec.match_file(candidate)
+
+
+def _priority(relative: Path) -> tuple:
+    suffix = relative.suffix.lower()
+    name = relative.name.lower()
+    depth = len(relative.parts)
+    if name in HIGH_PRIORITY_NAMES:
+        group = 0
+    elif relative.parts and relative.parts[0].lower() in {"test", "tests", "spec"}:
+        group = 3
+    elif suffix in {".py", ".ts", ".tsx", ".js", ".jsx", ".go", ".rs", ".java", ".cs"}:
+        group = 1
+    elif suffix in CODE_EXTS:
+        group = 2
+    else:
+        group = 3
+    return group, depth, relative.as_posix().lower()
+
+
+def _cache_path(root: Path) -> Path:
+    return root / ".zai" / "cache" / "repomap.json"
+
+
+def _load_cache(root: Path) -> dict:
+    data = read_json(_cache_path(root), {}, expected_type=dict, quarantine=True)
+    if data.get("version") != CACHE_VERSION:
+        return {"version": CACHE_VERSION, "files": {}}
+    return data
+
+
+def _cache_key(path: Path) -> str:
+    stat = path.stat()
+    return f"{stat.st_mtime_ns}:{stat.st_size}"
+
+
+def _content_hash(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _scan_candidates(
+    root: Path,
+    stats: RepoMapStats,
+    max_scan_files: int,
+    ignore_patterns: list[str] | None = None,
+    progress_callback=None,
+) -> list[tuple[Path, Path, int]]:
+    spec = _load_gitignore(root, ignore_patterns)
+    candidates = []
+    for current, dirs, files in os.walk(root, topdown=True, followlinks=False):
+        current_path = Path(current)
+        kept_dirs = []
+        for dirname in sorted(dirs):
+            path = current_path / dirname
+            relative = path.relative_to(root)
+            if path.is_symlink():
+                stats.symlinks += 1
+            elif dirname in SKIP_DIRS or _ignored(spec, relative.as_posix(), True):
+                stats.ignored += 1
+            else:
+                kept_dirs.append(dirname)
+        dirs[:] = kept_dirs
+
+        for filename in sorted(files):
+            if stats.scanned >= max_scan_files:
+                stats.truncated = True
+                return candidates
+            stats.scanned += 1
+            if progress_callback and stats.scanned % 250 == 0:
+                progress_callback(stats)
+            path = current_path / filename
+            relative = path.relative_to(root)
+            if path.is_symlink():
+                stats.symlinks += 1
+                continue
+            if _ignored(spec, relative.as_posix()) or _is_secret(relative):
+                stats.ignored += 1
+                continue
+            suffix = path.suffix.lower()
+            if suffix in SKIP_EXTS or suffix not in CODE_EXTS:
+                stats.unsupported += 1
+                continue
+            try:
+                size = path.stat().st_size
+            except OSError:
+                stats.errors += 1
+                continue
+            candidates.append((path, relative, size))
+    return candidates
+
+
+def build_repomap_result(
+    directory: str = ".",
+    max_files: int = DEFAULT_MAX_FILES,
+    *,
+    max_scan_files: int = DEFAULT_MAX_SCAN_FILES,
+    max_file_bytes: int = DEFAULT_MAX_FILE_BYTES,
+    max_repo_bytes: int = DEFAULT_MAX_REPO_BYTES,
+    use_cache: bool = True,
+    ignore_patterns: list[str] | None = None,
+    progress_callback=None,
+) -> RepoMapResult:
+    root = Path(directory).resolve()
+    if not root.is_dir():
+        raise ValueError(f"Repository directory not found: {directory}")
+    stats = RepoMapStats()
+    candidates = _scan_candidates(
+        root,
+        stats,
+        max_scan_files,
+        ignore_patterns,
+        progress_callback,
+    )
+    candidates.sort(key=lambda item: _priority(item[1]))
+    cache = _load_cache(root) if use_cache else {"version": CACHE_VERSION, "files": {}}
+    old_files = cache.get("files", {})
+    new_files = {}
+    entries = []
+    indexed_paths = []
+
+    for path, relative, size in candidates:
+        if len(entries) >= max_files:
+            stats.truncated = True
+            break
+        if size > max_file_bytes:
+            stats.oversized += 1
+            continue
+        if stats.bytes_considered + size > max_repo_bytes:
+            stats.truncated = True
+            break
+        if _looks_binary(path):
+            stats.binary += 1
+            continue
+        stats.bytes_considered += size
+        relative_text = relative.as_posix()
+        try:
+            key = _cache_key(path)
+        except OSError:
+            stats.errors += 1
+            continue
+        cached = old_files.get(relative_text)
+        if cached and cached.get("key") == key:
+            symbols = cached.get("symbols", [])
+            stats.cache_hits += 1
+        else:
+            symbols = _extract_symbols(path)
+            stats.cache_misses += 1
+        new_files[relative_text] = {
+            "key": key,
+            "sha256": (
+                cached.get("sha256")
+                if cached and cached.get("key") == key
+                else _content_hash(path)
+            ),
+            "symbols": symbols,
+            "size": size,
+        }
+        entries.append((relative_text, size, symbols))
+        indexed_paths.append(relative_text)
+
+    stats.indexed = len(entries)
+    if use_cache:
+        atomic_write_json(_cache_path(root), {
+            "version": CACHE_VERSION,
+            "root": str(root),
+            "files": new_files,
+        })
+
+    lines = [f"Repo: {root.name}", ""]
+    for relative, size, symbols in entries:
+        lines.append(f"{relative} ({size // 1024}kb)")
+        lines.extend(f"  {symbol}" for symbol in symbols)
+    lines.extend([
+        "",
+        f"Indexed: {stats.indexed} | Scanned: {stats.scanned} | "
+        f"Cache: {stats.cache_hits} hits/{stats.cache_misses} misses",
+        f"Skipped: ignored {stats.ignored}, unsupported {stats.unsupported}, "
+        f"oversized {stats.oversized}, binary {stats.binary}, "
+        f"symlinks {stats.symlinks}, errors {stats.errors}",
+    ])
+    if stats.truncated:
+        lines.append("Note: repository map was truncated by configured limits.")
+    return RepoMapResult("\n".join(lines), stats, indexed_paths)
+
+
+def build_repomap(directory: str = ".", max_files: int = DEFAULT_MAX_FILES) -> str:
+    return build_repomap_result(directory, max_files=max_files).text
+
+
+def get_repomap_prompt(
+    directory: str = ".",
+    *,
+    max_tokens: int = DEFAULT_PROMPT_TOKENS,
+) -> str:
+    result = build_repomap_result(directory)
+    prefix = "Here is the repository structure:\n\n"
+    suffix = "\n\nUse this map to understand the codebase when answering questions."
+    available_chars = max(512, (max_tokens - estimate_text_tokens(prefix + suffix)) * 3)
+    text = result.text
+    if len(text) > available_chars:
+        text = text[:available_chars] + "\n\n[Repository map truncated to token budget]"
+    return f"{prefix}{text}{suffix}"

zai/core/runtime.py

@@ -0,0 +1,29 @@
+"""Process-wide CLI presentation and diagnostics settings."""
+from __future__ import annotations
+
+from rich.console import Console
+
+_debug = False
+_plain = False
+
+
+def configure(*, debug: bool = False, plain: bool = False) -> None:
+    global _debug, _plain
+    _debug = debug
+    _plain = plain
+
+
+def debug_enabled() -> bool:
+    return _debug
+
+
+def plain_enabled() -> bool:
+    return _plain
+
+
+def print_exception(console: Console, error: Exception) -> None:
+    if _debug:
+        console.print_exception(show_locals=False)
+    else:
+        console.print(f"[red]Error:[/red] Unexpected error: {error}")
+        console.print("[dim]Run with --debug for full traceback.[/dim]")

zai/core/security.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+
+from .errors import FileError
+from .process import classify_argv, split_command
+
+
+def resolve_project_path(cwd: str | Path, user_path: str, *, allow_root: bool = False) -> Path:
+    """Resolve a user-controlled path and require it to stay inside the project."""
+    if not user_path or "\x00" in user_path:
+        raise FileError("Invalid empty path")
+
+    root = Path(cwd).resolve()
+    candidate = Path(user_path).expanduser()
+    if not candidate.is_absolute():
+        candidate = root / candidate
+    candidate = candidate.resolve()
+
+    try:
+        candidate.relative_to(root)
+    except ValueError as exc:
+        raise FileError(f"Path must stay inside the current project: {user_path}") from exc
+
+    if candidate == root and not allow_root:
+        raise FileError("Operation on the project root is not allowed")
+    return candidate
+
+
+def classify_command(command: str) -> tuple[str, str]:
+    """Return (safe|approval|blocked, reason) using direct argv policy."""
+    try:
+        return classify_argv(split_command(command))
+    except ValueError as exc:
+        return "blocked", str(exc)