vent_agent_2.0/runtime/python_sandbox.py

"""会话 Python 执行；工作区可写，代码快照与执行凭据由宿主保存。"""
import asyncio
from datetime import datetime
import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import signal
import sys
from zoneinfo import ZoneInfo

import config
from agents.workspace import session_root


def workspace_file(user, session, virtual):
    if not isinstance(virtual, str) or not virtual.startswith('/workspace/'):
        raise ValueError('文件路径必须以 /workspace/ 开头')
    base = session_root(user, session).resolve()
    path = (base / virtual.lstrip('/')).resolve()
    if not path.is_relative_to(base / 'workspace'):
        raise ValueError('路径越出本会话工作区')
    return path


def digest(path):
    value = hashlib.sha256()
    with path.open('rb') as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b''):
            value.update(chunk)
    return value.hexdigest()


def capability():
    if sys.platform == 'darwin' and shutil.which('sandbox-exec'):
        return {'available': True, 'backend': 'macOS sandbox-exec', 'network': False}
    if sys.platform.startswith('linux') and shutil.which('bwrap'):
        return {'available': True, 'backend': 'Linux bubblewrap', 'network': False}
    return {'available': False, 'reason': 'Python 执行需要 macOS sandbox-exec 或 Linux bubblewrap；不会回退为无沙箱执行'}


def _command(workspace, run_dir, script, timeout, memory):
    python = Path(sys.executable).absolute()
    worker = Path(__file__).with_name('python_worker.py').resolve()
    environment = {'PATH': os.defpath, 'HOME': str(workspace), 'TMPDIR': str(workspace / '.tmp'),
                   'LANG': 'en_US.UTF-8', 'PYTHONIOENCODING': 'utf-8',
                   'OPENBLAS_NUM_THREADS': '1', 'OMP_NUM_THREADS': '1',
                   'MPLCONFIGDIR': str(workspace / '.cache/matplotlib')}
    if sys.platform == 'darwin' and shutil.which('sandbox-exec'):
        readable = {workspace, run_dir, Path(sys.prefix).resolve(), Path(sys.base_prefix).resolve(),
                    Path('/System/Library'), Path('/usr/lib'), Path('/usr/share'), Path('/Library/Fonts')}
        quoted = lambda path: json.dumps(str(path), ensure_ascii=False)
        reads = ' '.join('(subpath ' + quoted(path) + ')' for path in sorted(readable))
        executables = {python.resolve()}
        # python.org 的 macOS Framework 解释器是启动器，会 re-exec 到 Python.app。
        framework_executable = Path(sys.base_prefix) / 'Resources/Python.app/Contents/MacOS/Python'
        if framework_executable.is_file():
            executables.add(framework_executable.resolve())
        execution_rules = ' '.join('(literal ' + quoted(path) + ')' for path in sorted(executables))
        # dyld 在解释器启动时会打开根目录定位共享缓存；仅放行目录本身，不递归放行宿主文件。
        profile = ('(version 1) (deny default) (allow file-read-metadata) '
                   '(allow file-read-data (literal "/")) '
                   '(allow sysctl-read) (allow signal (target self)) '
                   '(allow process-exec ' + execution_rules + ') '
                   '(allow file-read* ' + reads + ' (literal ' + quoted(worker) + ') '
                   '(literal "/dev/null") (literal "/dev/urandom") (literal "/dev/random")) '
                   '(allow file-map-executable ' + reads + ') '
                   '(allow file-write* (subpath ' + quoted(workspace) + ') (literal "/dev/null")) '
                   '(allow mach-lookup (global-name "com.apple.system.logger"))')
        # worker 本身不含配置、凭证；不向子进程传入服务端环境变量。
        return [shutil.which('sandbox-exec'), '-p', profile, str(python), '-I', '-B', '-u',
                str(worker), str(script), str(timeout), str(memory)], environment
    if sys.platform.startswith('linux') and shutil.which('bwrap'):
        command = [shutil.which('bwrap'), '--die-with-parent', '--new-session', '--unshare-all',
                   '--cap-drop', 'ALL', '--proc', '/proc', '--dev', '/dev', '--tmpfs', '/tmp']
        roots = {Path('/usr'), Path('/lib'), Path('/lib64'), Path(sys.prefix).resolve(), Path(sys.base_prefix).resolve()}
        for path in sorted(roots):
            if path.exists():
                command.extend(['--ro-bind', str(path), str(path)])
        command.extend(['--bind', str(workspace), '/workspace', '--ro-bind', str(run_dir), '/run-record',
                        '--ro-bind', str(worker), '/python_worker.py', '--chdir', '/workspace'])
        environment.update(HOME='/workspace', TMPDIR='/workspace/.tmp', MPLCONFIGDIR='/workspace/.cache/matplotlib')
        command.extend([str(python), '-I', '-B', '-u', '/python_worker.py',
                        '/run-record/' + script.name, str(timeout), str(memory)])
        return command, environment
    raise ValueError(capability()['reason'])


async def run(user, session, script_path, arguments=(), output_paths=(), timeout_seconds=60):
    script = workspace_file(user, session, script_path)
    if not script.is_file() or script.suffix != '.py' or script.stat().st_size > 1024 * 1024:
        raise ValueError('请选择工作区内不超过 1MB 的 Python 脚本')
    if len(arguments) > 32 or any(not isinstance(item, str) or len(item) > 4096 for item in arguments):
        raise ValueError('脚本参数数量或长度超过上限')
    if len(output_paths) > 20:
        raise ValueError('最多声明 20 个输出文件')
    outputs = [(virtual, workspace_file(user, session, virtual)) for virtual in output_paths]
    def output_version(path):
        if not path.is_file():
            return None
        state = path.stat()
        return state.st_ino, state.st_size, state.st_mtime_ns, state.st_ctime_ns
    prior_outputs = {virtual: output_version(path) for virtual, path in outputs}
    timeout = max(1, min(int(timeout_seconds), config.get_int('AGENT_PYTHON_MAX_SECONDS', 120)))
    root = session_root(user, session).resolve()
    workspace = root / 'workspace'
    workspace.mkdir(exist_ok=True)
    runs = root / 'python_runs'
    if runs.exists() and sum(1 for item in runs.iterdir() if item.is_dir()) >= config.get_int('AGENT_PYTHON_MAX_RUNS', 40):
        raise ValueError('本会话 Python 执行次数已达上限，请根据已有证据提交结果')
    (workspace / '.tmp').mkdir(exist_ok=True)
    identifier = os.urandom(16).hex()
    record_dir = root / 'python_runs' / identifier
    record_dir.mkdir(parents=True)
    snapshot = record_dir / 'script.py'
    snapshot.write_bytes(script.read_bytes())
    command, environment = _command(workspace, record_dir, snapshot, timeout,
                                    config.get_int('AGENT_PYTHON_MEMORY_MB', 2048))
    command.extend(arguments)
    limit = 64 * 1024

    async def consume(stream):
        buffer = bytearray()
        truncated = False
        while chunk := await stream.read(8192):
            remaining = max(0, limit - len(buffer))
            buffer.extend(chunk[:remaining])
            truncated |= len(chunk) > remaining
        return buffer.decode('utf-8', errors='replace'), truncated

    process = await asyncio.create_subprocess_exec(*command, cwd=workspace, env=environment,
        stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True)
    readers = [asyncio.create_task(consume(process.stdout)), asyncio.create_task(consume(process.stderr))]
    timed_out = False
    try:
        await asyncio.wait_for(process.wait(), timeout=timeout)
    except (TimeoutError, asyncio.CancelledError) as exc:
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
        await process.wait()
        if isinstance(exc, asyncio.CancelledError):
            await asyncio.gather(*readers, return_exceptions=True)
            raise
        timed_out = True
    finally:
        # 终止脚本留下的子进程，防止管道不关闭或任务结束后继续运行。
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            pass
    captured = await asyncio.gather(*readers)
    receipt = {'ok': process.returncode == 0 and not timed_out,
               'run_id': identifier, 'script_path': script_path, 'script_sha256': digest(snapshot),
               'exit_code': process.returncode, 'timed_out': timed_out,
               'stdout': captured[0][0], 'stderr': captured[1][0],
               'output_truncated': any(item[1] for item in captured), 'outputs': [],
               'declared_outputs': list(output_paths), 'unchanged_outputs': [], 'missing_outputs': [],
               'backend': capability().get('backend'), 'network': False,
               'created_at': datetime.now(ZoneInfo('Asia/Shanghai')).isoformat()}
    for virtual, path in outputs:
        # 脚本可能替换路径或建立符号链接，执行后重新解析边界。
        checked = workspace_file(user, session, virtual)
        if checked.is_file():
            if output_version(checked) == prior_outputs[virtual]:
                receipt['unchanged_outputs'].append(virtual)
            else:
                receipt['outputs'].append({'path': virtual, 'sha256': digest(checked), 'size': checked.stat().st_size})
        else:
            receipt['missing_outputs'].append(virtual)
    receipt['receipt_path'] = '/workspace/executions/' + identifier + '.json'
    encoded = json.dumps(receipt, ensure_ascii=False, indent=2)
    (record_dir / 'receipt.json').write_text(encoded, encoding='utf-8')
    readable = workspace_file(user, session, receipt['receipt_path'])
    readable.parent.mkdir(parents=True, exist_ok=True)
    readable.write_text(encoded, encoding='utf-8')
    return receipt


def execution_receipt(user, session, identifier):
    if not isinstance(identifier, str) or not re.fullmatch(r'[0-9a-f]{32}', identifier):
        raise ValueError('无效 Python 执行标识')
    path = session_root(user, session) / 'python_runs' / identifier / 'receipt.json'
    if not path.is_file():
        raise ValueError('Python 执行记录不存在或不属于当前会话')
    value = json.loads(path.read_text())
    if value['exit_code'] != 0 or value['timed_out']:
        raise ValueError('不能使用失败或超时的 Python 执行作为完成凭据')
    for output in value['outputs']:
        path = workspace_file(user, session, output['path'])
        if not path.is_file() or digest(path) != output['sha256']:
            raise ValueError('Python 输出文件已变化，请重新执行相应脚本')
    return value