cowork-local/core/deps.py

"""Runtime dependency helper.

Tasks and document extraction should never ask the user to install support
libraries by hand — when something is missing we try to ``pip install`` it into
the running interpreter automatically. In a packaged (frozen) build pip isn't
available, so callers must still degrade gracefully if this returns None/False.
"""
from __future__ import annotations

import importlib
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple

CancelFn = Callable[[], bool]

_POLL_SECS = 0.2

_FAILED: set[str] = set()  # packages we already tried and couldn't install

# 🔒 Sandbox Security Layer — Resource Usage: every subprocess spawned via
# run_cancellable registers its pid here for the duration of the run, so the
# Monitoring Dashboard can show live psutil stats without its own tracking.
_active_pids_lock = threading.Lock()
_ACTIVE_PIDS: set[int] = set()


def active_pids() -> List[int]:
    """Pid của các tiến trình con đang chạy — dùng để dọn sạch khi thoát app."""
    with _active_pids_lock:
        return sorted(_ACTIVE_PIDS)

# Substrings (lower-cased) in pip's output that mark a TRANSIENT failure (flaky
# network) worth silently retrying, as opposed to a deterministic one (bad
# package name, no matching version, syntax error in a requirement) where
# retrying would just waste time and reproduce the same error.
_TRANSIENT_MARKERS = (
    "connection reset", "connection aborted", "connection refused",
    "read timed out", "timed out", "temporary failure", "getaddrinfo failed",
    "could not fetch url", "network is unreachable", "max retries exceeded",
    "remote end closed connection", "econnreset",
)


def _kill_tree(proc: "subprocess.Popen", job_handle: Optional[int] = None) -> None:
    """Kill a subprocess AND any children it spawned (e.g. a shell wrapping the
    real command, or a build tool that forks workers) — plain ``proc.kill()``
    only kills the direct child and would leave the real work running.

    ``job_handle`` (Windows only), when the process was successfully assigned
    to one at spawn time, is tried FIRST — a Job Object catches reparented/
    detached processes that ``taskkill /T``'s PID-tree walk can miss (see
    win_job.py). Falls back to ``taskkill /T`` if there's no job handle."""
    if sys.platform == "win32" and job_handle:
        from .win_job import terminate_job

        if terminate_job(job_handle):
            try:
                proc.kill()
            except Exception:  # noqa: BLE001
                pass
            return
    try:
        if sys.platform == "win32":
            subprocess.run(
                ["taskkill", "/F", "/T", "/PID", str(proc.pid)],
                capture_output=True, timeout=10,
            )
        else:
            import os
            import signal
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            except (ProcessLookupError, PermissionError):
                pass
    except Exception:  # noqa: BLE001 - killing must never itself raise
        pass
    finally:
        try:
            proc.kill()
        except Exception:  # noqa: BLE001
            pass


def run_cancellable(
    args, *, cwd: str | None = None, timeout: float | None = None,
    cancel: Optional[CancelFn] = None, shell: bool = False,
    on_output: Optional[Callable[[str], None]] = None,
    env: Optional[Dict[str, str]] = None,
    limits: Optional[Dict[str, float]] = None,
) -> Tuple[Optional[int], str, bool, bool, bool]:
    """Run a subprocess so the Stop button can actually interrupt it.

    ``subprocess.run(..., timeout=...)`` blocks the calling thread until the
    process exits or the timeout fires — the cooperative cancel flag checked
    elsewhere in the agent loop has no chance to run, so Stop appears to do
    nothing while a command (or ``pip install``) is executing. This polls
    ``cancel()`` every ``_POLL_SECS`` instead and kills the whole process tree
    the moment the user stops, or the timeout is hit.

    ``on_output``, if given, is called with each line of stdout/stderr AS IT
    ARRIVES (not just at the end) so a caller can stream live progress to the
    UI for long-running commands — purely a side channel; the return value is
    unaffected.

    ``limits`` (Sandbox Security Layer — see ``resource_limits.py``), if
    given, is a dict of any subset of ``cpu_percent``/``memory_mb``/
    ``disk_mb``: the process TREE's usage is polled on the same cadence as
    cancel/timeout, and the tree is killed the moment a cap is exceeded.

    On Windows, the process is additionally assigned to a Job Object at spawn
    time (see ``win_job.py``) — a stronger tree-kill than ``taskkill /T``
    alone, since it also catches reparented/detached children. Best-effort:
    a failure to create/assign the job just means the existing taskkill
    fallback is used, same as before this was added.

    Returns ``(returncode, combined_output, cancelled, timed_out,
    resource_exceeded)``; on a failure to even launch the process,
    ``returncode`` is ``None`` and the output holds the launch error."""
    cancel = cancel or (lambda: False)
    popen_kwargs = {} if sys.platform == "win32" else {"start_new_session": True}
    try:
        proc = subprocess.Popen(
            args, shell=shell, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, bufsize=1, env=env, **popen_kwargs,
        )
    except OSError as exc:
        return None, str(exc), False, False, False

    with _active_pids_lock:
        _ACTIVE_PIDS.add(proc.pid)
    try:
        return _run_cancellable_body(proc, cancel, timeout, on_output, limits)
    finally:
        with _active_pids_lock:
            _ACTIVE_PIDS.discard(proc.pid)


def _run_cancellable_body(
    proc: "subprocess.Popen", cancel: CancelFn, timeout: Optional[float],
    on_output: Optional[Callable[[str], None]], limits: Optional[Dict[str, float]],
) -> Tuple[Optional[int], str, bool, bool, bool]:
    """Chạy một tiến trình con có thể huỷ giữa chừng, có hạn giờ và có giới hạn tài nguyên.

    Trên Windows gắn tiến trình vào một Job Object để khi giết là giết cả cây
    tiến trình con — giết mỗi tiến trình cha sẽ để lại đám con mồ côi.
    """
    job_handle = None
    if sys.platform == "win32":
        from .win_job import assign_process, create_job_object

        job_handle = create_job_object()
        if job_handle is not None:
            assign_process(job_handle, proc.pid)

    if limits:
        from .resource_limits import prime_cpu_counter

        prime_cpu_counter(proc.pid)

    collected: Dict[str, list] = {"out": [], "err": []}

    def _read_stream(stream, key: str) -> None:
        """Đọc một luồng đầu ra theo từng dòng ở luồng riêng.

        Phải đọc song song stdout và stderr: đọc lần lượt sẽ kẹt khi tiến trình con
        làm đầy bộ đệm của luồng còn lại.
        """
        try:
            for line in iter(stream.readline, ""):
                collected[key].append(line)
                if on_output is not None:
                    try:
                        on_output(line)
                    except Exception:  # noqa: BLE001 - a UI callback must never kill the tool
                        pass
        except Exception:  # noqa: BLE001
            pass
        finally:
            try:
                stream.close()
            except Exception:  # noqa: BLE001
                pass

    out_thread = threading.Thread(target=_read_stream, args=(proc.stdout, "out"), daemon=True)
    err_thread = threading.Thread(target=_read_stream, args=(proc.stderr, "err"), daemon=True)
    out_thread.start()
    err_thread.start()

    start = time.monotonic()
    cancelled = timed_out = resource_exceeded = False
    resource_reason = ""
    while out_thread.is_alive() or err_thread.is_alive():
        if cancel():
            cancelled = True
            _kill_tree(proc, job_handle)
            break
        if timeout is not None and (time.monotonic() - start) > timeout:
            timed_out = True
            _kill_tree(proc, job_handle)
            break
        if limits:
            from .resource_limits import check_limits

            resource_reason = check_limits(proc.pid, limits) or ""
            if resource_reason:
                resource_exceeded = True
                _kill_tree(proc, job_handle)
                break
        time.sleep(_POLL_SECS)
    out_thread.join(timeout=5)
    err_thread.join(timeout=5)
    try:
        proc.wait(timeout=5)  # reap so returncode is populated
    except subprocess.TimeoutExpired:
        pass

    out, err = "".join(collected["out"]), "".join(collected["err"])
    combined = out + (("\n[stderr]\n" + err) if err else "")
    if resource_exceeded:
        combined += f"\n[resource limit] {resource_reason}\n"
    return proc.returncode, combined, cancelled, timed_out, resource_exceeded


def network_blocked_env(base_env: Optional[Dict[str, str]] = None) -> Dict[str, str]:
    """Env vars that make well-behaved HTTP clients refuse to reach the
    network — proxy vars pointed at a black-hole loopback port nothing
    listens on (connection refused instantly, no hang).

    This is a POLICY-level control (Sandbox Security Layer — "Network
    Control"), not a kernel firewall: it stops the vast majority of scripted
    network calls (``requests``/``curl``/``wget``/``npm``/``pip`` all honor
    these standard proxy env vars) without requiring admin rights or a
    bundled driver — a tool that ignores proxy env vars entirely (rare, but
    possible) would still get through. Combine with the Agent Security
    command whitelist for defense in depth."""
    import os

    env = dict(base_env if base_env is not None else os.environ)
    blackhole = "http://127.0.0.1:1"
    for key in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"):
        env[key] = blackhole
    env["NO_PROXY"] = ""
    env["no_proxy"] = ""
    return env


def _can_pip() -> bool:
    # A PyInstaller/py2exe build has no usable pip; don't attempt installs there.
    """Bản đóng gói (PyInstaller) không có pip dùng được — đừng thử cài gì ở đó."""
    return not getattr(sys, "frozen", False)


def ensure_module(module: str, package: str | None = None):
    """Import ``module``, auto-installing ``package`` (pip) first if needed.

    Returns the imported module, or None if it isn't available and can't be
    installed (offline, no pip, frozen build, …)."""
    try:
        return importlib.import_module(module)
    except ImportError:
        pass
    pkg = package or module
    if pkg in _FAILED or not _can_pip():
        return None
    ok, _ = pip_install(pkg)
    if not ok:
        _FAILED.add(pkg)
        return None
    try:
        importlib.invalidate_caches()
        return importlib.import_module(module)
    except ImportError:
        _FAILED.add(pkg)
        return None


def venv_python_path(venv_dir: Path) -> Path:
    """Đường dẫn tới ``python`` trong một virtualenv, khác nhau giữa Windows và POSIX."""
    return venv_dir / ("Scripts/python.exe" if sys.platform == "win32" else "bin/python")


def ensure_project_venv(workdir: Path, cancel: Optional[CancelFn] = None,
                         on_output: Optional[Callable[[str], None]] = None) -> Optional[Path]:
    """Create (if missing) and return the interpreter of a per-project sandbox
    virtual environment at ``<workdir>/.venv`` — so packages the Code agent
    installs for one project never leak into another project's runs or into
    the app's own environment. Returns None (caller should fall back to the
    app's own interpreter) when a venv can't be created (offline, no pip, a
    packaged/frozen build, ...) — sandboxing is a nice-to-have, never a hard
    requirement for the agent to keep working."""
    if not _can_pip():
        return None
    venv_dir = workdir / ".venv"
    py = venv_python_path(venv_dir)
    if py.exists():
        return py
    if on_output is not None:
        on_output("[sandbox] creating project virtual environment (.venv)…\n")
    returncode, out, cancelled, _, _ = run_cancellable(
        [sys.executable, "-m", "venv", str(venv_dir)],
        timeout=120, cancel=cancel, on_output=on_output,
    )
    if returncode == 0 and py.exists():
        return py
    if on_output is not None:
        on_output(f"[sandbox] could not create .venv, using the app's own environment ({out.strip()[-300:]})\n")
    return None


def sandbox_env(python_path: Path) -> Dict[str, str]:
    """Env vars that make a subprocess behave as if this venv were activated —
    bare ``python``/``pip`` in a shell command then resolve to the sandbox."""
    import os

    env = dict(os.environ)
    bin_dir = str(Path(python_path).parent)
    env["PATH"] = bin_dir + os.pathsep + env.get("PATH", "")
    env["VIRTUAL_ENV"] = str(Path(python_path).parent.parent)
    env.pop("PYTHONHOME", None)
    return env


def pip_install(package: str, cancel: Optional[CancelFn] = None,
                 on_output: Optional[Callable[[str], None]] = None,
                 python: Optional[str] = None, retries: int = 2) -> tuple[bool, str]:
    """Install a pip package into ``python`` (default: the app's own interpreter).
    Returns (ok, output).

    Cancellable (see :func:`run_cancellable`) so hitting Stop while a package is
    installing actually kills pip instead of blocking until it finishes.

    A failure that looks like a flaky network blip (connection reset, timeout,
    DNS failure...) is retried automatically up to ``retries`` times with a
    short backoff — a deterministic failure (no matching version, bad package
    name) is NOT retried, since repeating it would just waste time."""
    if not _can_pip():
        return False, "This packaged build can't install packages at runtime."
    exe = python or sys.executable
    attempt = 0
    while True:
        attempt += 1
        returncode, out, cancelled, timed_out, _ = run_cancellable(
            [exe, "-m", "pip", "install", "--disable-pip-version-check", package],
            timeout=600, cancel=cancel, on_output=on_output,
        )
        if returncode is None:
            return False, f"pip failed to run: {out}"
        if cancelled:
            return False, "Installation cancelled by user."
        if timed_out:
            return False, "pip install timed out (600s) and was cancelled."
        if returncode == 0:
            return True, (out.strip()[-4000:] or "(no output)")
        transient = any(marker in out.lower() for marker in _TRANSIENT_MARKERS)
        if not transient or attempt > retries or (cancel and cancel()):
            return False, (out.strip()[-4000:] or "(no output)")
        if on_output is not None:
            on_output(f"\n[retry] transient network error — retrying ({attempt}/{retries})…\n")
        time.sleep(1.5 * attempt)