strix/strix/interface/tui/renderers/shell_renderer.py

import re
from functools import cache
from typing import Any, ClassVar

from pygments.lexers import get_lexer_by_name
from pygments.styles import get_style_by_name
from rich.text import Text
from textual.widgets import Static

from .base_renderer import BaseToolRenderer
from .registry import register_tool_renderer


MAX_OUTPUT_LINES = 50
MAX_LINE_LENGTH = 200

STRIP_PATTERNS = [
    r"^Chunk ID: [0-9a-f]+\s*$",
    r"^Wall time: [\d.]+ seconds\s*$",
    r"^Process exited with code -?\d+\s*$",
    r"^Process running with session ID \d+\s*$",
    r"^Original token count: \d+\s*$",
]

_EXIT_RE = re.compile(r"Process exited with code (-?\d+)")
_SESSION_RE = re.compile(r"Process running with session ID (\d+)")
_OUTPUT_HEADER = "\nOutput:\n"


@cache
def _get_style_colors() -> dict[Any, str]:
    style = get_style_by_name("native")
    return {token: f"#{style_def['color']}" for token, style_def in style if style_def["color"]}


def _parse_sdk_shell_result(result: Any) -> dict[str, Any]:
    """Translate the SDK's terminal-output string into the dict shape the
    renderer's `_append_output` helper expects.

    The SDK returns a header-prefixed string ending with `Output:\\n<actual>`.
    We extract `content`, `exit_code`, and `session_id`; anything else (or a
    non-string result) flows through unchanged so renderers can handle errors.
    """
    if isinstance(result, dict):
        return result
    if not isinstance(result, str):
        return {"content": "" if result is None else str(result)}

    exit_match = _EXIT_RE.search(result)
    session_match = _SESSION_RE.search(result)
    idx = result.find(_OUTPUT_HEADER)
    content = result[idx + len(_OUTPUT_HEADER) :] if idx >= 0 else result

    parsed: dict[str, Any] = {"content": content}
    if exit_match:
        parsed["exit_code"] = int(exit_match.group(1))
    if session_match:
        parsed["session_id"] = int(session_match.group(1))
    return parsed


def _truncate_line(line: str) -> str:
    clean_line = re.sub(r"\x1b\[[0-9;]*m", "", line)
    if len(clean_line) > MAX_LINE_LENGTH:
        return line[: MAX_LINE_LENGTH - 3] + "..."
    return line


def _clean_output(output: str) -> str:
    cleaned = output
    for pattern in STRIP_PATTERNS:
        cleaned = re.sub(pattern, "", cleaned, flags=re.MULTILINE)

    if cleaned.strip():
        lines = cleaned.splitlines()
        filtered_lines: list[str] = []
        for line in lines:
            if not filtered_lines and not line.strip():
                continue
            if line.strip() == "Output:":
                continue
            filtered_lines.append(line)
        while filtered_lines and not filtered_lines[-1].strip():
            filtered_lines.pop()
        cleaned = "\n".join(filtered_lines)

    return cleaned.strip()


def _format_output(output: str) -> Text:
    text = Text()
    lines = output.splitlines()
    total_lines = len(lines)

    head_count = MAX_OUTPUT_LINES // 2
    tail_count = MAX_OUTPUT_LINES - head_count - 1

    if total_lines <= MAX_OUTPUT_LINES:
        display_lines = lines
        truncated = False
        hidden_count = 0
    else:
        display_lines = lines[:head_count]
        truncated = True
        hidden_count = total_lines - head_count - tail_count

    for i, line in enumerate(display_lines):
        text.append("  ")
        text.append(_truncate_line(line), style="dim")
        if i < len(display_lines) - 1 or truncated:
            text.append("\n")

    if truncated:
        text.append(f"  ... {hidden_count} lines truncated ...", style="dim italic")
        text.append("\n")
        tail_lines = lines[-tail_count:]
        for i, line in enumerate(tail_lines):
            text.append("  ")
            text.append(_truncate_line(line), style="dim")
            if i < len(tail_lines) - 1:
                text.append("\n")

    return text


def _get_token_color(token_type: Any) -> str | None:
    colors = _get_style_colors()
    while token_type:
        if token_type in colors:
            return colors[token_type]
        token_type = token_type.parent
    return None


def _highlight_bash(code: str) -> Text:
    lexer = get_lexer_by_name("bash")
    text = Text()
    for token_type, token_value in lexer.get_tokens(code):
        if not token_value:
            continue
        color = _get_token_color(token_type)
        text.append(token_value, style=color)
    return text


def _append_output(text: Text, parsed: dict[str, Any], tool_status: str) -> None:
    raw_output = parsed.get("content", "") or ""
    output = _clean_output(raw_output) if isinstance(raw_output, str) else ""
    exit_code = parsed.get("exit_code")

    if tool_status == "running":
        if output:
            text.append("\n")
            text.append_text(_format_output(output))
        return

    if not output:
        if exit_code is not None and exit_code != 0:
            text.append("\n")
            text.append(f"  exit {exit_code}", style="dim #ef4444")
        return

    text.append("\n")
    text.append_text(_format_output(output))

    if exit_code is not None and exit_code != 0:
        text.append("\n")
        text.append(f"  exit {exit_code}", style="dim #ef4444")


def _build_terminal_content(
    *,
    prompt: str,
    prompt_style: str,
    command: str,
    parsed_result: dict[str, Any] | None,
    tool_status: str,
    meta: str | None = None,
) -> Text:
    text = Text()
    text.append(">_", style="dim")
    text.append(" ")

    if not command.strip():
        text.append("getting logs...", style="dim")
    else:
        text.append(prompt, style=prompt_style)
        text.append(" ")
        text.append_text(_highlight_bash(command))

    if meta:
        text.append(f"  {meta}", style="dim")

    if parsed_result is not None:
        _append_output(text, parsed_result, tool_status)

    return text


@register_tool_renderer
class ExecCommandRenderer(BaseToolRenderer):
    tool_name: ClassVar[str] = "exec_command"
    css_classes: ClassVar[list[str]] = ["tool-call", "terminal-tool"]

    @classmethod
    def render(cls, tool_data: dict[str, Any]) -> Static:
        args = tool_data.get("args", {})
        status = tool_data.get("status", "unknown")
        result = tool_data.get("result")

        cmd = str(args.get("cmd", ""))
        workdir = args.get("workdir")
        tty = bool(args.get("tty"))

        meta_parts: list[str] = []
        if workdir:
            meta_parts.append(f"cwd:{workdir}")
        if tty:
            meta_parts.append("tty")
        meta = ", ".join(meta_parts) if meta_parts else None

        parsed = _parse_sdk_shell_result(result) if result is not None else None

        content = _build_terminal_content(
            prompt="$",
            prompt_style="#22c55e",
            command=cmd,
            parsed_result=parsed,
            tool_status=status,
            meta=meta,
        )

        return Static(content, classes=cls.get_css_classes(status))


@register_tool_renderer
class WriteStdinRenderer(BaseToolRenderer):
    tool_name: ClassVar[str] = "write_stdin"
    css_classes: ClassVar[list[str]] = ["tool-call", "terminal-tool"]

    @classmethod
    def render(cls, tool_data: dict[str, Any]) -> Static:
        args = tool_data.get("args", {})
        status = tool_data.get("status", "unknown")
        result = tool_data.get("result")

        chars = str(args.get("chars", ""))
        session_id = args.get("session_id")
        meta = f"session #{session_id}" if session_id is not None else None

        parsed = _parse_sdk_shell_result(result) if result is not None else None

        content = _build_terminal_content(
            prompt=">>>",
            prompt_style="#3b82f6",
            command=chars,
            parsed_result=parsed,
            tool_status=status,
            meta=meta,
        )

        return Static(content, classes=cls.get_css_classes(status))