"""Remote target environment — transparently proxies tools to a GPU target.
The agent thinks it's working locally, but all operations execute on a remote
GPU target via SSH. This enables GPU-dependent operations (hipcc, profiling)
without the agent needing to know about remote execution.
Design:
- Wraps file operations: write/edit sync to remote, read pulls from remote
- Wraps bash: executes via SSH on remote target
- Maintains bidirectional sync of working directory
Usage:
    from wafer_core.environments.remote_target import RemoteTargetEnvironment
    env = RemoteTargetEnvironment(
        target_name="mi300x",
        working_dir=Path("/tmp/audit-workspace"),
    )
    # Agent tool calls are transparently proxied to remote target
    result = await env.exec_tool(tool_call, state, config)
"""
from __future__ import annotations
import logging
import os
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING
logger = logging.getLogger("wafer.eval.events")
import trio
from wafer_core.rollouts.dtypes import (
    AgentState,
    Message,
    RunConfig,
    Tool,
    ToolCall,
    ToolResult,
)
from wafer_core.tools import (
    BASH_TOOL,
    EDIT_TOOL,
    GLOB_TOOL,
    GREP_TOOL,
    READ_TOOL,
    WRITE_TOOL,
)
if TYPE_CHECKING:
    from wafer.targets_ops import TargetSSHInfo
# ══════════════════════════════════════════════════════════════════════════════
# Module-level SSH info cache for sharing across environments
# ══════════════════════════════════════════════════════════════════════════════
# This allows multiple RemoteTargetEnvironment instances (e.g., per-sample)
# to share the same provisioned target, avoiding re-provisioning for each sample.
_ssh_info_cache: dict[str, TargetSSHInfo] = {}
async def provision_target(target_name: str) -> TargetSSHInfo:
    """Provision a target upfront and cache the SSH info.
    Uses the proven kernelbench GPU provider infrastructure for provisioning.
    Args:
        target_name: Name of the target (from wafer target config)
    Returns:
        TargetSSHInfo for the provisioned target
    """
    if target_name in _ssh_info_cache:
        return _ssh_info_cache[target_name]
    from wafer.targets import load_target
    try:
        target = load_target(target_name)
    except FileNotFoundError:
        raise RuntimeError(
            f"Target '{target_name}' not found. "
            f"Configure with: wafer target config init {target_name} ..."
        ) from None

    target_type = type(target).__name__
    if target_type == "RunPodTarget":
        ssh_info = await _provision_runpod(target_name, target)
    elif target_type == "DigitalOceanTarget":
        ssh_info = await _provision_digitalocean(target_name, target)
    else:
        # Fall back to targets_ops for other target types (baremetal, etc.)
        from wafer.targets_ops import get_target_ssh_info
        print(f"Connecting to target '{target_name}'...")
        ssh_info = await get_target_ssh_info(target)
    _ssh_info_cache[target_name] = ssh_info
    print(f"Target '{target_name}' ready: {ssh_info.user}@{ssh_info.host}:{ssh_info.port}")
    return ssh_info
async def _provision_runpod(target_name: str, target: object) -> TargetSSHInfo:
    """Provision RunPod instance using kernelbench provider."""
    from wafer.targets_ops import TargetSSHInfo
    from wafer_core.auth import get_api_key
    # Try env var first, then wafer auth stored key
    api_key = os.environ.get("WAFER_RUNPOD_API_KEY") or get_api_key("runpod")
    if not api_key:
        raise RuntimeError(
            "RunPod API key required. Either:\n"
            "  - Set WAFER_RUNPOD_API_KEY environment variable, or\n"
            "  - Run: wafer auth login runpod"
        )
    from kernelbench.gpu_provider import get_provider
    print(f"Provisioning RunPod instance for '{target_name}'...")
    provider = get_provider("runpod", api_key=api_key)
    # Try to find existing instance first
    instance, is_new = await provider.get_or_provision(
        f"wafer-{target_name}",
        force_new=False,
    )
    # Always wait for ready to ensure we have direct SSH (not proxy)
    # RunPod proxy (ssh.runpod.io) doesn't handle non-interactive sessions well
    if is_new or instance.public_ip == "ssh.runpod.io":
        print("Waiting for direct SSH access (not proxy)...")
        ready = await provider.wait_for_ready(instance, timeout=600)
        if not ready:
            raise RuntimeError("Instance failed to become ready with direct SSH")
    ssh_key = getattr(target, "ssh_key", "~/.ssh/id_ed25519")
    expanded_key_path = Path(ssh_key).expanduser()  # noqa: ASYNC240 - sync string op
    return TargetSSHInfo(
        host=instance.public_ip or "",
        port=instance.ssh_port,
        user=instance.ssh_username,
        key_path=expanded_key_path,
    )
async def _provision_digitalocean(target_name: str, target: object) -> TargetSSHInfo:
    """Provision DigitalOcean instance using kernelbench provider."""
    from wafer.targets_ops import TargetSSHInfo
    from wafer_core.auth import get_api_key
    # Try env var first, then wafer auth stored key
    api_key = os.environ.get("WAFER_AMD_DIGITALOCEAN_API_KEY") or get_api_key("digitalocean")
    if not api_key:
        raise RuntimeError(
            "DigitalOcean API key required. Either:\n"
            "  - Set WAFER_AMD_DIGITALOCEAN_API_KEY environment variable, or\n"
            "  - Run: wafer auth login digitalocean"
        )
    from kernelbench.gpu_provider import get_provider
    print(f"Provisioning DigitalOcean instance for '{target_name}'...")
    provider = get_provider("digitalocean", api_key=api_key)
    # Try to find existing instance first
    instance, is_new = await provider.get_or_provision(
        f"wafer-{target_name}",
        force_new=False,
    )
    # Always wait for ready on new instances
    if is_new:
        print("Waiting for instance to be ready...")
        ready = await provider.wait_for_ready(instance, timeout=600)
        if not ready:
            raise RuntimeError("Instance failed to become ready")
    ssh_key = getattr(target, "ssh_key", "~/.ssh/id_ed25519")
    expanded_key_path = Path(ssh_key).expanduser()  # noqa: ASYNC240 - sync string op
    return TargetSSHInfo(
        host=instance.public_ip or "",
        port=instance.ssh_port,
        user=instance.ssh_username,
        key_path=expanded_key_path,
    )
def clear_ssh_cache() -> None:
    """Clear the SSH info cache. Useful for testing."""
    _ssh_info_cache.clear()
# Tool definitions available in remote environment
ALL_TOOLS = {
    "read": READ_TOOL,
    "write": WRITE_TOOL,
    "edit": EDIT_TOOL,
    "glob": GLOB_TOOL,
    "grep": GREP_TOOL,
    "bash": BASH_TOOL,
}
@dataclass
class RemoteTargetEnvironment:
    """Environment that transparently proxies all operations to a remote GPU target.
    Args:
        target_name: Name of the target (from wafer target config).
        working_dir: Local working directory. Files are synced to remote.
        remote_working_dir: Remote working directory (default: /tmp/wafer-workspace).
        enabled_tools: List of tool names to enable. If None, all tools enabled.
        bash_timeout: Timeout for bash commands in seconds.
        sync_on_init: Whether to sync working_dir to remote on initialization.
    """
    target_name: str
    working_dir: Path = field(default_factory=Path.cwd)
    remote_working_dir: str = "/tmp/wafer-workspace"
    enabled_tools: list[str] | None = None
    bash_timeout: int = 300
    sync_on_init: bool = True
    # Internal state
    _ssh_info: TargetSSHInfo | None = field(init=False, repr=False, default=None)
    _synced_files: set[str] = field(init=False, repr=False, default_factory=set)

    _ready: bool = field(init=False, repr=False, default=False)

    _initial_sync_done: bool = field(init=False, repr=False, default=False)
    def __post_init__(self) -> None:
        """Validate configuration."""
        if self.enabled_tools is not None:
            valid_tools = set(ALL_TOOLS.keys())
            unknown = set(self.enabled_tools) - valid_tools
            if unknown:
                raise ValueError(
                    f"Unknown tools: {sorted(unknown)}. Available: {sorted(valid_tools)}"
                )
    async def ensure_ready(self) -> None:
        """Provision target and sync files upfront.
        Call this before starting the eval to avoid provisioning delays
        during the first tool call. Idempotent — safe to call multiple times.
        """
        if self._ready:
            return
        await self._ensure_ssh_info()
        self._ready = True
    async def _ensure_ssh_info(self) -> TargetSSHInfo:
        """Get SSH info for target, using cache if available."""
        if self._ssh_info is not None:
            return self._ssh_info
        if self.target_name in _ssh_info_cache:
            self._ssh_info = _ssh_info_cache[self.target_name]
        else:
            # Provision and cache
            self._ssh_info = await provision_target(self.target_name)
        logger.debug(
            "Creating remote directory",
            extra={
                "target": self.target_name,
                "remote_dir": self.remote_working_dir,
                "host": self._ssh_info.host,
            },
        )
        await self._exec_remote_raw(f"mkdir -p {self.remote_working_dir}", self._ssh_info)
        # Initial sync: do it once if enabled, then mark as done (survives serialize/deserialize)
        if self.sync_on_init and not self._initial_sync_done and self.working_dir.exists():
            sync_start = time.time()
            file_count = len(list(self.working_dir.iterdir()))
            logger.info(
                "Starting initial sync to remote",
                extra={
                    "target": self.target_name,
                    "local_dir": str(self.working_dir),
                    "remote_dir": self.remote_working_dir,
                    "file_count": file_count,
                    "sync_on_init": self.sync_on_init,
                    "initial_sync_done": self._initial_sync_done,
                },
            )
            await self._sync_to_remote()
            self._initial_sync_done = True
            logger.info(
                "Initial sync complete",
                extra={
                    "target": self.target_name,
                    "duration_ms": int((time.time() - sync_start) * 1000),
                    "file_count": file_count,
                },
            )
        return self._ssh_info
    async def _exec_remote_raw(
        self,
        command: str,
        ssh_info: TargetSSHInfo,
    ) -> tuple[int, str, str]:
        """Execute command on remote without cd wrapper (for bootstrap).
        Use this only for commands that need to run before the remote working
        directory exists (e.g., mkdir -p to create it).
        Callers that need a timeout should wrap with trio.move_on_after().
        """
        wrapped_command = f"BASH_ENV= bash -c '{command}'"
        ssh_args = [
            "ssh",
            "-T",
            "-i", str(ssh_info.key_path),
            "-p", str(ssh_info.port),
            "-o", "StrictHostKeyChecking=no",
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", "LogLevel=ERROR",
            "-o", "BatchMode=yes",
            "-o", "RequestTTY=no",
            f"{ssh_info.user}@{ssh_info.host}",
            wrapped_command,
        ]
        result = await trio.run_process(
            ssh_args,
            capture_stdout=True,
            capture_stderr=True,
            check=False,
        )
        return (
            result.returncode,
            result.stdout.decode() if isinstance(result.stdout, bytes) else result.stdout,
            result.stderr.decode() if isinstance(result.stderr, bytes) else result.stderr,
        )
    async def _exec_remote(
        self,
        command: str,
    ) -> tuple[int, str, str]:
        """Execute command on remote target.
        Callers that need a timeout should wrap with trio.move_on_after().
        Returns:
            (exit_code, stdout, stderr)
        """
        ssh_info = await self._ensure_ssh_info()
        # Wrap command with BASH_ENV unset to skip .bashrc which may check for PTY
        wrapped_command = f"BASH_ENV= bash -c 'cd {self.remote_working_dir} && {command}'"
        ssh_args = [
            "ssh",
            "-T",  # Disable PTY allocation
            "-i",
            str(ssh_info.key_path),
            "-p",
            str(ssh_info.port),
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            "-o",
            "LogLevel=ERROR",
            "-o",
            "BatchMode=yes",  # Non-interactive mode
            "-o",
            "RequestTTY=no",  # Explicitly request no TTY
            f"{ssh_info.user}@{ssh_info.host}",
            wrapped_command,
        ]
        exec_start = time.time()
        result = await trio.run_process(
            ssh_args,
            capture_stdout=True,
            capture_stderr=True,
            check=False,
        )
        exit_code = result.returncode
        stdout = result.stdout.decode() if isinstance(result.stdout, bytes) else result.stdout
        stderr = result.stderr.decode() if isinstance(result.stderr, bytes) else result.stderr
        duration_ms = int((time.time() - exec_start) * 1000)
        # Wide event: one log per remote execution with full context
        logger.debug(
            "remote_exec",
            extra={
                "target": self.target_name,
                "host": ssh_info.host,
                "command": command[:200],  # Truncate long commands
                "exit_code": exit_code,
                "duration_ms": duration_ms,
                "stdout_len": len(stdout) if stdout else 0,
                "stderr_len": len(stderr) if stderr else 0,
                "success": exit_code == 0,
            },
        )
        if exit_code != 0:
            logger.warning(
                "Remote command failed",
                extra={
                    "target": self.target_name,
                    "command": command[:200],
                    "exit_code": exit_code,
                    "stderr": stderr[:500] if stderr else None,
                },
            )
        return (exit_code, stdout, stderr)
    async def _sync_to_remote(self, paths: list[Path] | None = None) -> None:
        """Sync files from local to remote using scp.
        Args:
            paths: Specific paths to sync. If None, syncs entire working_dir.
        """
        ssh_info = await self._ensure_ssh_info()
        scp_base = [
            "scp",
            "-r",  # Recursive for directories
            "-T",  # Disable PTY
            "-i", str(ssh_info.key_path),
            "-P", str(ssh_info.port),
            "-o", "StrictHostKeyChecking=no",
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", "LogLevel=ERROR",
            "-o", "BatchMode=yes",
        ]
        if paths is None:
            # Sync entire working directory contents
            sources = list(self.working_dir.iterdir())
            if not sources:
                logger.debug(
                    "No files to sync",
                    extra={"target": self.target_name, "working_dir": str(self.working_dir)},
                )
                return
            scp_args = scp_base + [str(s) for s in sources]
            scp_args.append(f"{ssh_info.user}@{ssh_info.host}:{self.remote_working_dir}/")
            num_items = len(sources)
        else:
            # Sync specific files
            scp_args = scp_base + [str(p.resolve()) for p in paths]
            scp_args.append(f"{ssh_info.user}@{ssh_info.host}:{self.remote_working_dir}/")
            num_items = len(paths)
        def run_scp() -> tuple[int, str, str]:
            result = subprocess.run(scp_args, capture_output=True, text=True)
            return (result.returncode, result.stdout, result.stderr)
        scp_start = time.time()
        exit_code, stdout, stderr = await trio.to_thread.run_sync(run_scp)
        duration_ms = int((time.time() - scp_start) * 1000)
        # Wide event: one log per scp operation with full context
        logger.info(
            "scp_sync",
            extra={
                "target": self.target_name,
                "host": ssh_info.host,
                "port": ssh_info.port,
                "num_items": num_items,
                "remote_dir": self.remote_working_dir,
                "exit_code": exit_code,
                "duration_ms": duration_ms,
                "success": exit_code == 0,
                "stderr": stderr[:500] if stderr else None,  # Truncate for large errors
            },
        )
        if exit_code != 0:
            logger.error(
                "scp sync failed",
                extra={
                    "target": self.target_name,
                    "exit_code": exit_code,
                    "stderr": stderr,
                    "scp_command": " ".join(scp_args[:8]) + "...",  # First few args
                },
            )
    async def _sync_from_remote(self, remote_path: str, local_path: Path) -> None:
        """Sync a file from remote to local."""
        ssh_info = await self._ensure_ssh_info()
        scp_args = [
            "scp",
            "-T",  # Disable PTY
            "-i", str(ssh_info.key_path),
            "-P", str(ssh_info.port),
            "-o", "StrictHostKeyChecking=no",
            "-o", "UserKnownHostsFile=/dev/null",
            "-o", "LogLevel=ERROR",
            "-o", "BatchMode=yes",
            f"{ssh_info.user}@{ssh_info.host}:{remote_path}",
            str(local_path),
        ]
        def run_scp() -> None:
            subprocess.run(scp_args, capture_output=True, check=True)
        await trio.to_thread.run_sync(run_scp)
    def get_name(self) -> str:
        """Return environment name identifier."""
        return f"remote-target:{self.target_name}"
    async def serialize(self) -> dict:
        return {
            "target_name": self.target_name,
            "working_dir": str(self.working_dir),
            "remote_working_dir": self.remote_working_dir,
            "enabled_tools": self.enabled_tools,
            "bash_timeout": self.bash_timeout,
            "_initial_sync_done": self._initial_sync_done,
        }
    @staticmethod
    async def deserialize(data: dict) -> RemoteTargetEnvironment:
        env = RemoteTargetEnvironment(
            target_name=data["target_name"],
            working_dir=Path(data["working_dir"]),
            remote_working_dir=data.get("remote_working_dir", "/tmp/wafer-workspace"),
            enabled_tools=data.get("enabled_tools"),
            bash_timeout=data.get("bash_timeout", 300),
        )
        # Restore sync state so we don't re-sync on every deserialize
        env._initial_sync_done = data.get("_initial_sync_done", False)
        return env
    def requires_confirmation(self, tool_call: ToolCall) -> bool:
        """Remote bash commands don't need local confirmation (sandboxed by SSH)."""
        return False
    def get_tools(self) -> list[Tool]:
        """Return enabled tools."""
        if self.enabled_tools is None:
            return list(ALL_TOOLS.values())
        return [ALL_TOOLS[name] for name in self.enabled_tools if name in ALL_TOOLS]
    async def on_assistant_message(self, message: Message, state: AgentState) -> AgentState:
        """No feedback needed."""
        return state
    async def exec_tool(
        self,
        tool_call: ToolCall,
        current_state: AgentState,
        run_config: RunConfig,
        cancel_scope: trio.CancelScope | None = None,
    ) -> ToolResult:
        """Execute tool call, proxying to remote target."""
        if self.enabled_tools is not None and tool_call.name not in self.enabled_tools:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"Tool '{tool_call.name}' is not enabled.",
            )
        try:
            if tool_call.name == "bash":
                return await self._exec_bash(tool_call)
            elif tool_call.name == "read":
                return await self._exec_read(tool_call)
            elif tool_call.name == "write":
                return await self._exec_write(tool_call)
            elif tool_call.name == "edit":
                return await self._exec_edit(tool_call)
            elif tool_call.name == "glob":
                return await self._exec_glob(tool_call)
            elif tool_call.name == "grep":
                return await self._exec_grep(tool_call)
            else:
                return ToolResult(
                    tool_call_id=tool_call.id,
                    is_error=True,
                    content="",
                    error=f"Unknown tool: {tool_call.name}",
                )
        except Exception as e:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=str(e),
            )
    async def _exec_bash(self, tool_call: ToolCall) -> ToolResult:
        """Execute bash command on remote target.
        Special handling for `wafer tool isa analyze`: runs locally after syncing
        the .co file from remote, since wafer CLI needs API auth.
        """
        args = tool_call.args
        command = args.get("command", "")
        timeout = args.get("timeout", self.bash_timeout * 1000) // 1000  # ms -> s
        # Intercept wafer tool isa analyze commands - run locally with synced files
        if "wafer tool isa" in command and "analyze" in command:
            return await self._exec_wafer_isa_analyze(tool_call, command)
        exit_code, stdout, stderr = (-1, "", f"Command timed out after {timeout}s")
        with trio.move_on_after(timeout):
            exit_code, stdout, stderr = await self._exec_remote(command)
        # Format output like local bash tool
        output = stdout
        if stderr:
            output = f"{stdout}\n{stderr}" if stdout else stderr
        if exit_code != 0:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content=output,
                error=f"Exit code {exit_code}",
            )
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=output,
        )
    async def _exec_wafer_isa_analyze(self, tool_call: ToolCall, command: str) -> ToolResult:
        """Handle `wafer tool isa analyze` by syncing files and running locally.
        The wafer CLI needs API auth which is configured locally, so we:
        1. Parse the file path from the command
        2. Sync the .co file from remote to local temp dir
        3. Run wafer tool isa analyze locally
        4. Return the output
        """
        import re
        import tempfile
        # or "cd kernels && wafer tool isa analyze foo.co"
        match = re.search(r"wafer\s+tool\s+isa\s+analyze\s+([^\s]+)", command)
        if not match:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error="Could not parse file path from wafer tool isa analyze command",
            )
        file_arg = match.group(1)
        cd_match = re.search(r"cd\s+([^\s&;]+)", command)
        if cd_match and not file_arg.startswith("/"):
            file_arg = f"{cd_match.group(1)}/{file_arg}"
        remote_path = self._to_remote_path(file_arg)
        suffix = Path(file_arg).suffix or ".co"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            local_path = Path(tmp.name)
        try:
            # Sync file from remote
            logger.info(
                "Syncing file for local ISA analysis",
                extra={
                    "target": self.target_name,
                    "remote_path": remote_path,
                    "local_path": str(local_path),
                },
            )
            await self._sync_from_remote(remote_path, local_path)
            # Run wafer tool isa analyze locally
            extra_flags = ""
            if "--json" in command or "-j" in command:
                extra_flags += " --json"
            wafer_cmd = f"wafer tool isa analyze {local_path}{extra_flags}"
            def run_local() -> tuple[int, str, str]:
                result = subprocess.run(
                    wafer_cmd,
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=60,
                )
                return (result.returncode, result.stdout, result.stderr)
            exit_code, stdout, stderr = await trio.to_thread.run_sync(run_local)
            output = stdout
            if stderr:
                output = f"{stdout}\n{stderr}" if stdout else stderr
            logger.info(
                "Local ISA analysis complete",
                extra={
                    "target": self.target_name,
                    "file": file_arg,
                    "exit_code": exit_code,
                    "output_len": len(output),
                },
            )
            if exit_code != 0:
                return ToolResult(
                    tool_call_id=tool_call.id,
                    is_error=True,
                    content=output,
                    error=f"Exit code {exit_code}",
                )
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=False,
                content=output,
            )
        finally:
            # Cleanup temp file
            local_path.unlink(missing_ok=True)
    async def _exec_read(self, tool_call: ToolCall) -> ToolResult:
        """Read file from remote target."""
        args = tool_call.args
        file_path = args.get("path", "")
        offset = args.get("offset", 0)
        limit = args.get("limit")
        remote_path = self._to_remote_path(file_path)
        if limit:
            cmd = f"cat -n '{remote_path}' | tail -n +{offset + 1} | head -n {limit}"
        elif offset:
            cmd = f"cat -n '{remote_path}' | tail -n +{offset + 1}"
        else:
            cmd = f"cat -n '{remote_path}'"
        exit_code, stdout, stderr = (-1, "", f"Command timed out after {self.bash_timeout}s")
        with trio.move_on_after(self.bash_timeout):
            exit_code, stdout, stderr = await self._exec_remote(cmd)
        if exit_code != 0:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=stderr or f"Failed to read {file_path}",
            )
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=stdout,
        )
    async def _exec_write(self, tool_call: ToolCall) -> ToolResult:
        """Write file locally and sync to remote."""
        args = tool_call.args
        file_path = args.get("path", "")
        content = args.get("content", "")
        # Resolve path
        local_path = self._resolve_local_path(file_path)
        # Write locally
        local_path.parent.mkdir(parents=True, exist_ok=True)
        local_path.write_text(content)
        # Sync to remote
        await self._sync_to_remote([local_path])
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=f"Wrote {len(content)} bytes to {file_path}",
        )
    async def _exec_edit(self, tool_call: ToolCall) -> ToolResult:
        """Edit file: pull from remote, edit locally, push back."""
        args = tool_call.args
        file_path = args.get("path", "")
        old_string = args.get("old_string", "")
        new_string = args.get("new_string", "")
        replace_all = args.get("replace_all", False)
        # Resolve paths
        local_path = self._resolve_local_path(file_path)
        remote_path = self._to_remote_path(file_path)
        # Pull latest from remote
        local_path.parent.mkdir(parents=True, exist_ok=True)
        try:
            await self._sync_from_remote(remote_path, local_path)
        except Exception:
            # File might not exist on remote yet
            if not local_path.exists():
                return ToolResult(
                    tool_call_id=tool_call.id,
                    is_error=True,
                    content="",
                    error=f"File not found: {file_path}",
                )
        # Read content
        content = local_path.read_text()
        if old_string not in content:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"old_string not found in {file_path}",
            )
        if not replace_all and content.count(old_string) > 1:
            return ToolResult(
                tool_call_id=tool_call.id,
                is_error=True,
                content="",
                error=f"old_string appears {content.count(old_string)} times. Use replace_all=true or provide more context.",
            )
        # Replace
        if replace_all:
            new_content = content.replace(old_string, new_string)
        else:
            new_content = content.replace(old_string, new_string, 1)
        # Write locally
        local_path.write_text(new_content)
        # Sync to remote
        await self._sync_to_remote([local_path])
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=f"Edited {file_path}",
        )
    async def _exec_glob(self, tool_call: ToolCall) -> ToolResult:
        """Execute glob on remote target."""
        args = tool_call.args
        pattern = args.get("pattern", "")
        path = args.get("path", ".")
        cmd = f"find {path} -name '{pattern}' -type f 2>/dev/null | head -100"
        exit_code, stdout, stderr = (-1, "", f"Command timed out after {self.bash_timeout}s")
        with trio.move_on_after(self.bash_timeout):
            exit_code, stdout, stderr = await self._exec_remote(cmd)
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=stdout if stdout else "No files found",
        )
    async def _exec_grep(self, tool_call: ToolCall) -> ToolResult:
        """Execute grep on remote target."""
        args = tool_call.args
        pattern = args.get("pattern", "")
        path = args.get("path", ".")
        output_mode = args.get("output_mode", "files_with_matches")
        if output_mode == "files_with_matches":
            cmd = f"grep -rl '{pattern}' {path} 2>/dev/null | head -50"
        elif output_mode == "count":
            cmd = f"grep -rc '{pattern}' {path} 2>/dev/null | grep -v ':0$' | head -50"
        else:  # content
            context = args.get("-C", 0)
            if context:
                cmd = f"grep -rn -C {context} '{pattern}' {path} 2>/dev/null | head -200"
            else:
                cmd = f"grep -rn '{pattern}' {path} 2>/dev/null | head -200"
        exit_code, stdout, stderr = (-1, "", f"Command timed out after {self.bash_timeout}s")
        with trio.move_on_after(self.bash_timeout):
            exit_code, stdout, stderr = await self._exec_remote(cmd)
        return ToolResult(
            tool_call_id=tool_call.id,
            is_error=False,
            content=stdout if stdout else "No matches found",
        )
    def _resolve_local_path(self, file_path: str) -> Path:
        """Resolve file path to local filesystem."""
        p = Path(file_path).expanduser()
        if not p.is_absolute():
            p = self.working_dir / p
        return p.resolve()
    def _to_remote_path(self, file_path: str) -> str:
        """Convert local path to remote path."""
        p = Path(file_path).expanduser()
        if p.is_absolute():
            # Try to make relative to working_dir
            try:
                rel = p.resolve().relative_to(self.working_dir.resolve())
                return f"{self.remote_working_dir}/{rel}"
            except ValueError:
                # Path not under working_dir, use as-is
                return str(p)
        else:
            return f"{self.remote_working_dir}/{file_path}"
