tmux.py

import asyncio  # noqa -- swapping to trio would be beneficial, but not blocking atm
import os
import re
import shutil
import time
from datetime import datetime, UTC

from .base import CLIResult, ToolError, ToolResult

DEFAULT_SESSION_NAME = "base"
DEFAULT_TIMEOUT_S = 10.0
DEFAULT_FRAME_RATE_HZ = 5.0
NEW_SESSION_DELAY_S = 2.0
MAX_OUTPUT_CHARS = 8 * 1024
_TMUX_COMMAND_TIMEOUT_S = 30.0


def _truncate_from_middle(value: str | None, max_chars: int = MAX_OUTPUT_CHARS) -> str | None:
    """Keep the head and tail of long output, dropping the middle."""
    if value is None or len(value) <= max_chars:
        return value
    half = max_chars // 2
    omitted = len(value) - 2 * half
    return f"{value[:half]}\n... [{omitted} characters truncated] ...\n{value[-half:]}"


def _utc_timestamp() -> str:
    """Millisecond UTC timestamp, same format as the hosted tool's per-frame timestamps."""
    return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"


def _strip_empty_lines(value: str) -> str:
    value = re.sub(r"^(\s*\n)+", "", value)
    value = re.sub(r"(\n\s*)+$", "", value)
    return value


def _remove_overlap(previous: str | None, current: str) -> str:
    """Drop lines of `current` already shown at the end of `previous` (hosted-tool frame dedup).

    The last line of the previous frame is always repeated, since it commonly keeps changing
    (a prompt gaining typed text, a spinner, a progress bar).
    """
    if not previous or not current:
        return current
    previous_lines = previous.splitlines()[0:-1]
    current_lines = current.splitlines()
    max_overlap_lines = 0
    for i in range(min(len(previous_lines), len(current_lines)), 1, -1):
        if previous_lines[-i:] == current_lines[:i]:
            max_overlap_lines = i
            break
    return "\n".join(current_lines[max_overlap_lines:])


class TmuxTool:
    """Lightweight local stand-in for the tmux tool provided by hosted Taiga.

    Interface mirrors the hosted tool:

    - `args` are passed directly to the tmux binary (`["send-keys", "-t", "base", "ls", "Enter"]`).
    - A default detached session named "base" is created on first use.
    - Passing a `-t <target>` with anything other than a direct `capture-pane` call triggers an
      automatic capture loop: the pane is polled at `frame_rate` until its content stops changing,
      one of the `patterns` regexes matches, or `timeout` seconds elapse. Changed frames are
      returned as a sequence of timestamped `<header>`/`<footer>` blocks with consecutive frames
      deduped — the same output format the hosted tool produces.
    - Multiline text goes in the `text` parameter with a literal `$text` placeholder in `args`.
    - `restart=True` kills the tmux server and recreates the default session.

    If the environment performs well with this tool, it will also work with the internal
    implementation used on hosted Taiga.

    The hosted implementation lives in the anthropic monorepo: `model_tools/model_tools/tmux_tool.py`
    (interface, defaults, warnings) and `model_tools_core/model_tools_core/line_io_sandbox_tools/_tmux_tool.py`
    (capture loop, frame format, dedup). Model-visible behavior changes there should be mirrored here.
    """

    def __init__(self):
        self._default_session_started = False
        # Last stripped frame per capture target, so frames dedupe across calls like the
        # hosted tool (whose capture server keeps the same state for the lifetime of the tool).
        # Keyed by the raw target string — not a resolved pane id — also like the hosted tool:
        # "base" and "base:0.0" are deliberately tracked as separate targets.
        self._previous_capture_by_target: dict[str, str] = {}

    async def __call__(
        self,
        *,
        args: list[str] | None = None,
        text: str | None = None,
        timeout: float | None = None,
        frame_rate: float | None = None,
        patterns: list[str] | None = None,
        restart: bool = False,
    ) -> ToolResult:
        if shutil.which("tmux") is None:
            raise ToolError(
                "tmux is not installed in this container. Add it to your Dockerfile "
                "(e.g. `apt-get install -y tmux`) to use the tmux tool."
            )

        if restart:
            await self._run_tmux(["kill-server"])
            self._default_session_started = False
            self._previous_capture_by_target.clear()
            if not args:
                await self._ensure_default_session()
                return ToolResult(system="tool has been restarted.")

        if not args:
            raise ToolError("no args provided.")

        if text is None and "$text" in args:
            raise ToolError('When using "$text" in your args, you must provide `text` as well')
        if text is not None and "$text" not in args:
            raise ToolError('When using the `text` parameter you must provide "$text" in your args as well')
        if text is not None:
            args = [text if arg == "$text" else arg for arg in args]

        args = [str(arg) for arg in args]
        capture_params_given = any(value is not None for value in (timeout, frame_rate, patterns))

        # Compile patterns eagerly so an invalid regex fails cleanly before any tmux command
        # runs or capture state is touched. re.DOTALL matches the hosted tool, so multi-line
        # patterns behave the same locally.
        compiled_patterns: list[re.Pattern[str]] | None = None
        if patterns is not None:
            try:
                compiled_patterns = [re.compile(pattern, re.DOTALL) for pattern in patterns]
            except re.error as e:
                raise ToolError(f"invalid regex in `patterns`: {e}") from e

        warnings: list[str] = []
        if capture_params_given and "-t" not in args:
            warnings.append(
                '`timeout`, `frame_rate`, and `patterns` are ignored when "-t" is not provided within your `args`'
            )
        if args[0] == "capture-pane":
            if capture_params_given:
                warnings.append(
                    "`timeout`, `frame_rate`, and `patterns` are ignored when you directly call capture-pane"
                )
            if "-p" not in args:
                warnings.append(
                    'You will not be able to see the result of capture-pane unless you also pass "-p" within your `args`'
                )

        timeout = DEFAULT_TIMEOUT_S if timeout is None else timeout
        frame_rate = DEFAULT_FRAME_RATE_HZ if frame_rate is None else frame_rate

        result = await self._ensure_default_session()
        result = result + await self._run_tmux(args)
        if warnings:
            result = result + ToolResult(system="\n".join(warnings))

        if args[0] == "capture-pane":
            return self._finalize(result)

        if args[0] == "new-session":
            await asyncio.sleep(NEW_SESSION_DELAY_S)

        target = self._extract_target(args)
        if target is None:
            return self._finalize(result)

        capture = await self._capture_until_settled(
            target=target, timeout=timeout, frame_rate=frame_rate, patterns=compiled_patterns
        )
        return self._finalize(result + capture)

    def _finalize(self, result: ToolResult) -> ToolResult:
        return result.replace(
            output=_truncate_from_middle(result.output),
            error=_truncate_from_middle(result.error),
        )

    @staticmethod
    def _extract_target(args: list[str]) -> str | None:
        try:
            return args[args.index("-t") + 1]
        except (ValueError, IndexError):
            return None

    async def _ensure_default_session(self) -> ToolResult:
        """Create the default detached session on first use."""
        if self._default_session_started:
            return ToolResult()
        # No explicit -x/-y: the hosted tool uses tmux's default pane size, and pane width
        # affects line wrapping, capture-pane output, and `patterns` matching.
        result = await self._run_tmux(["new-session", "-d", "-s", DEFAULT_SESSION_NAME])
        self._default_session_started = True
        if result.error and "duplicate session" in result.error:
            # The session is already there (e.g. a previous tool instance created it) — that's fine.
            return ToolResult()
        return result

    async def _capture_until_settled(
        self,
        *,
        target: str,
        timeout: float,
        frame_rate: float,
        patterns: list[re.Pattern[str]] | None,
    ) -> ToolResult:
        """Poll capture-pane and return the sequence of changed frames, like the hosted tool.

        Each changed frame is wrapped in timestamped `<header>`/`<footer>` blocks and deduped
        against the previous frame. The loop stops when the pane is static, a pattern matches
        (an `<info>` block is appended), or the timeout elapses (a `<warning>` block is appended).
        """
        deadline = time.monotonic() + timeout
        interval = 1.0 / frame_rate if frame_rate > 0 else 1.0 / DEFAULT_FRAME_RATE_HZ
        collected: list[str] = []
        frame_index = 0
        while True:
            now = time.monotonic()
            if now >= deadline:
                # The hosted tool is cancelled exactly at the deadline; emit the same warning.
                collected.append(f"<warning>Timeout of {timeout}s was reached</warning>")
                break
            next_frame_at = now + interval

            timestamp_before = _utc_timestamp()
            # -e (keep escape sequences) and -J (join wrapped lines) match the hosted tool's
            # capture flags — both change the text that `patterns` are matched against.
            frame = await self._run_tmux(["capture-pane", "-e", "-p", "-J", "-t", target])
            timestamp_after = _utc_timestamp()
            # The hosted tool merges capture stderr into the frame content, so errors
            # (e.g. a bad target) stay visible to the model.
            content = _strip_empty_lines((frame.output or "") + (frame.error or ""))

            previous = self._previous_capture_by_target.get(target)
            self._previous_capture_by_target[target] = content
            has_changed = content != previous
            deduped = _remove_overlap(previous, content)

            # The first frame is always shown, even when unchanged since the previous call.
            if has_changed or frame_index == 0:
                collected.append(
                    f"<header><frame>{frame_index}</frame>"
                    f"<timestamp>{timestamp_before}</timestamp><target>{target}</target></header>"
                )
                collected.append(deduped)
                collected.append(f"\033[0m<footer><timestamp>{timestamp_after}</timestamp></footer>")

            settled = False
            if patterns is None:
                # Require at least two frames so a pane that happens to match the previous
                # call's final state doesn't end the capture before the command has run.
                settled = frame_index != 0 and not has_changed
            else:
                for pattern in patterns:
                    # Patterns are matched against the deduped frame, like the hosted tool.
                    match = pattern.search(deduped)
                    if match:
                        collected.append(
                            f"<info>frame={frame_index} matched <pattern>{pattern.pattern}</pattern>"
                            f" with <match>{match.group(0)}</match></info>"
                        )
                        settled = True
            if settled:
                break

            frame_index += 1
            # Sleep until the next frame slot, but never past the deadline — the hosted tool
            # hard-stops at `timeout` even at low frame rates.
            await asyncio.sleep(max(0.0, min(next_frame_at, deadline) - time.monotonic()))
        return ToolResult(output="\n".join(collected))

    async def _run_tmux(self, args: list[str]) -> CLIResult:
        """Run a single tmux command (no shell) as the model user and collect its output."""

        def demote():
            # This only runs in the child process — same pattern as BashTool.
            os.setsid()
            os.setgid(1000)
            os.setuid(1000)

        process = await asyncio.create_subprocess_exec(
            "tmux",
            *args,
            preexec_fn=demote,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=_TMUX_COMMAND_TIMEOUT_S)
        except TimeoutError:
            process.kill()
            return CLIResult(error=f"tmux {' '.join(args)} timed out after {_TMUX_COMMAND_TIMEOUT_S} seconds")
        return CLIResult(
            output=stdout.decode(errors="replace"),
            error=stderr.decode(errors="replace"),
        )