OpenOPC/opc/layer4_tools/shell.py

"""Structured shell execution tools."""

from __future__ import annotations

import asyncio
import os
import shutil
from pathlib import Path
from typing import Any, AsyncIterator

from opc.layer4_tools.execution_context import (
    build_subprocess_env,
    resolve_task_execution_context,
    wrap_command_for_context,
)
from opc.layer4_tools.registry import ToolDefinition
from opc.layer2_organization.work_item_identity import work_item_turn_type_from_metadata


_STDOUT_LIMIT = 50_000
_STDERR_LIMIT = 20_000
_SETUP_STAGE_DEFAULT_TIMEOUT = 1800
_DEFAULT_SHELL_TIMEOUT = 300
_POWERSHELL_CMD_SEPARATOR = " ; "
_BASH_CMD_SEPARATOR = " && "
_STREAM_READ_SIZE = 8192


def _resolve_working_directory(
    working_directory: str | None = None,
    task: Any | None = None,
) -> str:
    cwd = str(working_directory or "").strip()
    if not cwd and task is not None:
        metadata = getattr(task, "metadata", {}) or {}
        execution_context = dict(metadata.get("_execution_context", {}) or {})
        candidates = [
            str(execution_context.get("workspace_root", "") or "").strip(),
            str(execution_context.get("output_root", "") or "").strip(),
            str(metadata.get("workspace_root", "") or "").strip(),
            str(metadata.get("comms_workspace_root", "") or "").strip(),
            str(metadata.get("output_root", "") or "").strip(),
            str(metadata.get("target_output_dir", "") or "").strip(),
        ]
        fallback = ""
        for raw in candidates:
            if not raw:
                continue
            path = Path(raw).expanduser()
            if path.exists() and path.is_dir():
                return str(path)
            if not fallback:
                fallback = str(path)
        cwd = fallback
    return cwd or os.getcwd()


def _shell_binary(preferred: str, fallback: str) -> str:
    return shutil.which(preferred) or fallback


async def _run_shell_command(
    *,
    shell_name: str,
    command: str,
    args: list[str],
    working_directory: str | None = None,
    timeout: int = _DEFAULT_SHELL_TIMEOUT,
    task: Any | None = None,
    on_progress: Any = None,
) -> dict[str, Any]:
    cwd = _resolve_working_directory(working_directory, task)
    cwd_path = Path(cwd).expanduser()
    resolved_cwd = str(cwd_path.resolve(strict=False))
    if not cwd_path.exists() or not cwd_path.is_dir():
        return {
            "success": False,
            "shell": shell_name,
            "command": command,
            "cwd": resolved_cwd,
            "stdout": "",
            "stderr": "",
            "exit_code": -1,
            "timed_out": False,
            "error": f"Working directory does not exist: {resolved_cwd}",
            "sandbox": {
                "platform": "",
                "requested_mode": "",
                "effective_mode": "off",
                "available": False,
                "fallback_used": False,
            },
            "execution_context": {},
        }
    if task is not None:
        meta = getattr(task, "metadata", {}) or {}
        override = meta.get("shell_timeout_override")
        if override is not None:
            try:
                timeout = max(int(override), timeout)
            except (ValueError, TypeError):
                pass
        elif work_item_turn_type_from_metadata(meta, fallback="") == "setup":
            timeout = max(timeout, _SETUP_STAGE_DEFAULT_TIMEOUT)
        shell_prefix = ""
        shell_prefix_win = ""
        inherited = meta.get("inherited_environment")
        if isinstance(inherited, dict):
            shell_prefix = str(inherited.get("shell_prefix", "") or "").strip()
            shell_prefix_win = str(inherited.get("shell_prefix_win", "") or "").strip()
        if not shell_prefix:
            manifest = meta.get("environment_manifest")
            if isinstance(manifest, dict):
                shell_prefix = str(manifest.get("shell_prefix", "") or "").strip()
                shell_prefix_win = str(manifest.get("shell_prefix_win", "") or "").strip()
        is_powershell = shell_name == "powershell"
        active_prefix = shell_prefix_win if (is_powershell and shell_prefix_win) else shell_prefix
        if active_prefix and active_prefix not in command:
            separator = _POWERSHELL_CMD_SEPARATOR if is_powershell else _BASH_CMD_SEPARATOR
            command = f"{active_prefix}{separator}{command}"
            # Replace only the trailing command argument. PowerShell args are
            # [exe, "-NoProfile", "-Command", command] (4 elements); the previous
            # ``[args[0], args[1], command]`` dropped the "-Command" flag and left
            # PowerShell unable to interpret the prefixed command. Bash args are
            # [bash, "-lc", command] (3 elements), handled identically here.
            args = [*args[:-1], command] if len(args) >= 1 else args
    context = resolve_task_execution_context(task)
    if resolved_cwd and not context.get("workspace_root"):
        context["workspace_root"] = resolved_cwd
    env = build_subprocess_env(context)
    try:
        wrapped_args, sandbox_meta = wrap_command_for_context(args, cwd=resolved_cwd, context=context)
    except RuntimeError as exc:
        return {
            "success": False,
            "shell": shell_name,
            "command": command,
            "cwd": resolved_cwd,
            "stdout": "",
            "stderr": "",
            "exit_code": -1,
            "timed_out": False,
            "error": str(exc),
            "sandbox": {
                "platform": (context.get("sandbox", {}) or {}).get("platform", ""),
                "requested_mode": (context.get("sandbox", {}) or {}).get("mode", ""),
                "effective_mode": "off",
                "available": False,
                "fallback_used": False,
            },
            "execution_context": _context_preview(context),
        }
    proc: asyncio.subprocess.Process | None = None
    try:
        proc = await asyncio.create_subprocess_exec(
            *wrapped_args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=resolved_cwd,
            env=env,
        )
        stdout_chunks: list[str] = []
        stderr_chunks: list[str] = []

        async def _pump(stream: Any, bucket: list[str], stream_name: str, limit: int) -> None:
            async for chunk in _iter_stream_lines(stream):
                text = chunk.decode("utf-8", errors="replace")
                bucket.append(text)
                joined = "".join(bucket)
                if len(joined) > limit:
                    bucket[:] = [joined[:limit]]
                if on_progress:
                    try:
                        await on_progress(text.rstrip("\r\n"), stream=stream_name)
                    except TypeError:
                        await on_progress(text.rstrip("\r\n"))

        stdout_task = asyncio.create_task(_pump(proc.stdout, stdout_chunks, "stdout", _STDOUT_LIMIT))
        stderr_task = asyncio.create_task(_pump(proc.stderr, stderr_chunks, "stderr", _STDERR_LIMIT))
        await asyncio.wait_for(proc.wait(), timeout=timeout)
        await asyncio.gather(stdout_task, stderr_task)
        stdout = "".join(stdout_chunks)[:_STDOUT_LIMIT]
        stderr = "".join(stderr_chunks)[:_STDERR_LIMIT]
        return {
            "success": proc.returncode == 0,
            "shell": shell_name,
            "command": command,
            "cwd": resolved_cwd,
            "stdout": stdout,
            "stderr": stderr,
            "exit_code": proc.returncode,
            "timed_out": False,
            "sandbox": sandbox_meta,
            "execution_context": _context_preview(context),
        }
    except asyncio.TimeoutError:
        if proc is not None:
            proc.kill()
        return {
            "success": False,
            "shell": shell_name,
            "command": command,
            "cwd": resolved_cwd,
            "stdout": "",
            "stderr": "",
            "exit_code": -1,
            "timed_out": True,
            "error": f"Command timed out after {timeout}s",
            "sandbox": sandbox_meta,
            "execution_context": _context_preview(context),
        }


async def _iter_stream_lines(stream: asyncio.StreamReader) -> AsyncIterator[bytes]:
    buffer = bytearray()
    while True:
        chunk = await stream.read(_STREAM_READ_SIZE)
        if not chunk:
            if buffer:
                yield bytes(buffer)
            return
        buffer.extend(chunk)
        while True:
            newline_index = buffer.find(b"\n")
            if newline_index < 0:
                break
            line = bytes(buffer[: newline_index + 1])
            del buffer[: newline_index + 1]
            yield line


def _context_preview(context: dict[str, Any]) -> dict[str, Any]:
    sandbox = dict(context.get("sandbox", {}) or {})
    return {
        "workspace_root": str(context.get("workspace_root", "") or ""),
        "output_root": str(context.get("output_root", "") or ""),
        "comms_root": str(context.get("comms_root", "") or ""),
        "venv_path": str(context.get("venv_path", "") or ""),
        "python_executable": str(context.get("python_executable", "") or ""),
        "venv_provider": str(context.get("venv_provider", "") or ""),
        "preparation_error": str(context.get("preparation_error", "") or ""),
        "sandbox": sandbox,
    }


async def bash_exec(
    command: str,
    working_directory: str | None = None,
    timeout: int = _DEFAULT_SHELL_TIMEOUT,
    task: Any | None = None,
    on_progress: Any = None,
) -> dict[str, Any]:
    """Execute a command using bash/sh semantics."""
    shell_binary = _shell_binary("bash", "sh" if os.name != "nt" else "bash")
    return await _run_shell_command(
        shell_name="bash",
        command=command,
        args=[shell_binary, "-lc", command],
        working_directory=working_directory,
        timeout=timeout,
        task=task,
        on_progress=on_progress,
    )


async def powershell_exec(
    command: str,
    working_directory: str | None = None,
    timeout: int = _DEFAULT_SHELL_TIMEOUT,
    task: Any | None = None,
    on_progress: Any = None,
) -> dict[str, Any]:
    """Execute a command using PowerShell semantics."""
    executable = shutil.which("pwsh") or shutil.which("powershell") or "powershell"
    return await _run_shell_command(
        shell_name="powershell",
        command=command,
        args=[executable, "-NoProfile", "-Command", command],
        working_directory=working_directory,
        timeout=timeout,
        task=task,
        on_progress=on_progress,
    )


async def shell_exec(
    command: str,
    working_directory: str | None = None,
    timeout: int = _DEFAULT_SHELL_TIMEOUT,
    shell: str | None = None,
    task: Any | None = None,
    on_progress: Any = None,
) -> dict[str, Any]:
    """Compatibility wrapper that selects bash or PowerShell."""
    normalized = str(shell or "").strip().lower()
    if normalized == "powershell":
        return await powershell_exec(
            command=command,
            working_directory=working_directory,
            timeout=timeout,
            task=task,
            on_progress=on_progress,
        )
    if os.name == "nt" and normalized not in {"bash", "sh"}:
        return await powershell_exec(
            command=command,
            working_directory=working_directory,
            timeout=timeout,
            task=task,
            on_progress=on_progress,
        )
    return await bash_exec(
        command=command,
        working_directory=working_directory,
        timeout=timeout,
        task=task,
        on_progress=on_progress,
    )


def _shell_schema(description: str) -> dict[str, Any]:
    return {
        "type": "object",
        "properties": {
            "command": {"type": "string", "description": description},
            "working_directory": {"type": "string", "description": "Working directory for the command (optional)"},
            "timeout": {"type": "integer", "description": "Timeout in seconds", "default": _DEFAULT_SHELL_TIMEOUT},
        },
        "required": ["command"],
    }


def create_shell_tools() -> list[ToolDefinition]:
    return [
        ToolDefinition(
            name="shell_exec",
            description="Execute a shell command. Selects bash or PowerShell based on platform or the optional `shell` hint.",
            parameters={
                "type": "object",
                "properties": {
                    **_shell_schema("The shell command to execute")["properties"],
                    "shell": {
                        "type": "string",
                        "description": "Optional shell hint: bash | powershell",
                        "default": "",
                    },
                },
                "required": ["command"],
            },
            func=shell_exec,
            category="compute",
            requires_confirmation=True,
            concurrency_safe=False,
            read_only=False,
        ),
    ]


def create_shell_tool() -> ToolDefinition:
    """Backward-compatible helper used by older callers."""
    for tool in create_shell_tools():
        if tool.name == "shell_exec":
            return tool
    raise RuntimeError("shell_exec definition is missing")