openclaw-skill/skills/code-interpreter/scripts/run_code.py

#!/usr/bin/env python3
import argparse
import importlib.util
import json
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile
import time
from typing import Optional

WORKSPACE = pathlib.Path('/home/selig/.openclaw/workspace').resolve()
RUNS_DIR = WORKSPACE / '.tmp' / 'code-interpreter-runs'
MAX_PREVIEW = 12000
ARTIFACT_SCAN_LIMIT = 100
PACKAGE_PROBES = ['pandas', 'numpy', 'matplotlib']
PYTHON_BIN = str(WORKSPACE / '.venv-code-interpreter' / 'bin' / 'python')


def current_python_paths(run_dir_path: pathlib.Path) -> str:
    """Build PYTHONPATH: run_dir (for ci_helpers) only.
    Venv site-packages are already on sys.path when using PYTHON_BIN."""
    return str(run_dir_path)


def read_code(args: argparse.Namespace) -> str:
    sources = [bool(args.code), bool(args.file), bool(args.stdin)]
    if sum(sources) != 1:
        raise SystemExit('Provide exactly one of --code, --file, or --stdin')
    if args.code:
        return args.code
    if args.file:
        return pathlib.Path(args.file).read_text(encoding='utf-8')
    return sys.stdin.read()


def ensure_within_workspace(path_str: Optional[str], must_exist: bool = True) -> pathlib.Path:
    if not path_str:
        return WORKSPACE
    p = pathlib.Path(path_str).expanduser().resolve()
    if p != WORKSPACE and WORKSPACE not in p.parents:
        raise SystemExit(f'Path must stay inside workspace: {WORKSPACE}')
    if must_exist and (not p.exists() or not p.is_dir()):
        raise SystemExit(f'Path not found or not a directory: {p}')
    return p


def ensure_output_path(path_str: Optional[str]) -> Optional[pathlib.Path]:
    if not path_str:
        return None
    p = pathlib.Path(path_str).expanduser().resolve()
    p.parent.mkdir(parents=True, exist_ok=True)
    return p


def write_text(path_str: Optional[str], text: str) -> None:
    p = ensure_output_path(path_str)
    if not p:
        return
    p.write_text(text, encoding='utf-8')


def truncate(text: str) -> str:
    if len(text) <= MAX_PREVIEW:
        return text
    extra = len(text) - MAX_PREVIEW
    return text[:MAX_PREVIEW] + f'\n...[truncated {extra} chars]'


def package_status() -> dict:
    out: dict[str, bool] = {}
    for name in PACKAGE_PROBES:
        proc = subprocess.run(
            [PYTHON_BIN, '-c', f"import importlib.util; print('1' if importlib.util.find_spec('{name}') else '0')"],
            capture_output=True,
            text=True,
            encoding='utf-8',
            errors='replace',
        )
        out[name] = proc.stdout.strip() == '1'
    return out


def rel_to(path: pathlib.Path, base: pathlib.Path) -> str:
    try:
        return str(path.relative_to(base))
    except Exception:
        return str(path)


def scan_artifacts(base_dir: pathlib.Path, root_label: str) -> list[dict]:
    if not base_dir.exists():
        return []
    items: list[dict] = []
    for p in sorted(base_dir.rglob('*')):
        if len(items) >= ARTIFACT_SCAN_LIMIT:
            break
        if p.is_file():
            try:
                size = p.stat().st_size
            except Exception:
                size = None
            items.append({
                'root': root_label,
                'path': str(p),
                'relative': rel_to(p, base_dir),
                'bytes': size,
            })
    return items


def write_helper(run_dir_path: pathlib.Path, artifact_dir: pathlib.Path) -> None:
    helper = run_dir_path / 'ci_helpers.py'
    helper.write_text(
        """
from pathlib import Path
import json
import os

WORKSPACE = Path(os.environ['OPENCLAW_WORKSPACE'])
RUN_DIR = Path(os.environ['CODE_INTERPRETER_RUN_DIR'])
ARTIFACT_DIR = Path(os.environ['CODE_INTERPRETER_ARTIFACT_DIR'])


def save_text(name: str, text: str) -> str:
    path = ARTIFACT_DIR / name
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(text, encoding='utf-8')
    return str(path)


def save_json(name: str, data) -> str:
    path = ARTIFACT_DIR / name
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
    return str(path)
""".lstrip(),
        encoding='utf-8',
    )


def main() -> int:
    parser = argparse.ArgumentParser(description='Local Python runner for OpenClaw code-interpreter skill')
    parser.add_argument('--code', help='Python code to execute')
    parser.add_argument('--file', help='Path to a Python file to execute')
    parser.add_argument('--stdin', action='store_true', help='Read Python code from stdin')
    parser.add_argument('--cwd', help='Working directory inside workspace')
    parser.add_argument('--artifact-dir', help='Artifact directory inside workspace to keep outputs')
    parser.add_argument('--timeout', type=int, default=20, help='Timeout seconds (default: 20)')
    parser.add_argument('--stdout-file', help='Optional file path to save full stdout')
    parser.add_argument('--stderr-file', help='Optional file path to save full stderr')
    parser.add_argument('--keep-run-dir', action='store_true', help='Keep generated temp run directory even on success')
    args = parser.parse_args()

    code = read_code(args)
    cwd = ensure_within_workspace(args.cwd)
    RUNS_DIR.mkdir(parents=True, exist_ok=True)

    run_dir_path = pathlib.Path(tempfile.mkdtemp(prefix='run-', dir=str(RUNS_DIR))).resolve()
    artifact_dir = ensure_within_workspace(args.artifact_dir, must_exist=False) if args.artifact_dir else (run_dir_path / 'artifacts')
    artifact_dir.mkdir(parents=True, exist_ok=True)

    script_path = run_dir_path / 'main.py'
    script_path.write_text(code, encoding='utf-8')
    write_helper(run_dir_path, artifact_dir)

    env = {
        'PATH': os.environ.get('PATH', '/usr/bin:/bin'),
        'HOME': str(run_dir_path),
        'PYTHONPATH': current_python_paths(run_dir_path),
        'PYTHONIOENCODING': 'utf-8',
        'PYTHONUNBUFFERED': '1',
        'OPENCLAW_WORKSPACE': str(WORKSPACE),
        'CODE_INTERPRETER_RUN_DIR': str(run_dir_path),
        'CODE_INTERPRETER_ARTIFACT_DIR': str(artifact_dir),
        'MPLBACKEND': 'Agg',
    }

    started = time.time()
    timed_out = False
    exit_code = None
    stdout = ''
    stderr = ''

    try:
        proc = subprocess.run(
            [PYTHON_BIN, '-B', str(script_path)],
            cwd=str(cwd),
            env=env,
            capture_output=True,
            text=True,
            encoding='utf-8',
            errors='replace',
            timeout=max(1, args.timeout),
        )
        exit_code = proc.returncode
        stdout = proc.stdout
        stderr = proc.stderr
    except subprocess.TimeoutExpired as exc:
        timed_out = True
        exit_code = 124
        raw_out = exc.stdout or ''
        raw_err = exc.stderr or ''
        stdout = raw_out if isinstance(raw_out, str) else raw_out.decode('utf-8', errors='replace')
        stderr = (raw_err if isinstance(raw_err, str) else raw_err.decode('utf-8', errors='replace')) + f'\nExecution timed out after {args.timeout}s.'

    duration = round(time.time() - started, 3)

    write_text(args.stdout_file, stdout)
    write_text(args.stderr_file, stderr)

    artifacts = scan_artifacts(artifact_dir, 'artifactDir')
    if artifact_dir != run_dir_path:
        artifacts.extend(scan_artifacts(run_dir_path / 'artifacts', 'runArtifacts'))

    result = {
        'ok': (exit_code == 0 and not timed_out),
        'exitCode': exit_code,
        'timeout': timed_out,
        'durationSec': duration,
        'cwd': str(cwd),
        'runDir': str(run_dir_path),
        'artifactDir': str(artifact_dir),
        'packageStatus': package_status(),
        'artifacts': artifacts,
        'stdout': truncate(stdout),
        'stderr': truncate(stderr),
    }

    print(json.dumps(result, ensure_ascii=False, indent=2))

    if not args.keep_run_dir and result['ok'] and artifact_dir != run_dir_path:
        shutil.rmtree(run_dir_path, ignore_errors=True)

    return 0 if result['ok'] else 1


if __name__ == '__main__':
    raise SystemExit(main())