This commit is contained in:
thanhnv
+350
@@ -0,0 +1,350 @@
|
||||
"""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]:
|
||||
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]:
|
||||
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:
|
||||
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.
|
||||
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:
|
||||
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)
|
||||
Reference in New Issue
Block a user