statuslinepy-sub

#!/usr/bin/env python3
# pyright: strict
"""Render the Claude Code subagent status line.

Rich owns terminal styling while Humanize supplies SI token prefixes for task
token counts. Accepts subagent task batches on stdin and outputs JSON task rows.

Format: model effort | elapsed | tokens | task
"""

from __future__ import annotations

import json
import os
import re
import sys
import time
from dataclasses import dataclass
from decimal import ROUND_HALF_UP, Decimal, InvalidOperation
from typing import TypeAlias, cast

import humanize
from rich.cells import cell_len
from rich.color import ColorSystem
from rich.style import Style

BLUE = "rgb(0,153,255)"
ORANGE = "rgb(255,176,85)"
GREEN = "rgb(0,160,0)"
RED = "rgb(255,85,85)"
PURPLE = "rgb(167,139,250)"
WHITE = "rgb(220,220,220)"
DIM = "dim"

MAX_INTEGER = 999_999_999_999_999
MAX_COLUMNS = 10_000
MIN_EPOCH_MILLISECONDS = 1_000_000_000_000
MAX_EPOCH_MILLISECONDS = 9_999_999_999_999
ZERO_DECIMAL = Decimal(0)
MAX_TOKEN_DISPLAY = Decimal("999000000")
MODEL_INPUT_CAP = 256
TASK_INPUT_CAP = 256
EFFORT_INPUT_CAP = 64
MODEL_WIDTH_CAP = 32
EFFORT_WIDTH_CAP = 8
ELAPSED_WIDTH_CAP = 10
TOKEN_WIDTH_CAP = 6
TASK_WIDTH_CAP = 256

JSONScalar: TypeAlias = str | int | float | bool | None
JSONValue: TypeAlias = (
    JSONScalar | list["JSONValue"] | dict[str, "JSONValue"]
)
JSONObject: TypeAlias = dict[str, JSONValue]


class Text:
    """Accumulate literal bytes-to-be with program-owned Rich styles.

    Rich's high-level Text renderer normalizes terminal controls embedded in
    content. The Bash contract preserves those bytes, so literal content stays
    in raw segments and only Style.render participates in final serialization.
    """

    def __init__(self, value: object = "", style: str | None = None) -> None:
        self.segments: list[tuple[str, str | None]] = []
        content = str(value)
        if content or style is not None:
            self.segments.append((content, style))

    @property
    def plain(self) -> str:
        return "".join(content for content, _style in self.segments)

    def append_text(self, other: Text) -> None:
        self.segments.extend(other.segments)

    def append(self, value: str) -> None:
        if value:
            self.segments.append((value, None))

    def __bool__(self) -> bool:
        return bool(self.segments)

    def render(self) -> str:
        rendered: list[str] = []
        for content, style_name in self.segments:
            if style_name is None:
                rendered.append(content)
                continue
            style = Style.parse(style_name)
            if content:
                rendered.append(
                    style.render(content, color_system=ColorSystem.TRUECOLOR)
                )
            else:
                marker = style.render("x", color_system=ColorSystem.TRUECOLOR)
                rendered.append(marker.replace("x", "", 1))
        return "".join(rendered)


def text(value: object = "", style: str | None = None) -> Text:
    """Return literal text with only a program-owned Rich style applied."""
    return Text(str(value), style=style)


def combine(*parts: Text | str) -> Text:
    result = Text()
    for part in parts:
        if isinstance(part, Text):
            result.append_text(part)
        else:
            result.append(part)
    return result


def visible_len(value: Text) -> int:
    """Return terminal display-cell width of text, stripping SGR codes."""
    plain = re.sub(r"\x1b\[[0-9;]*m", "", value.plain)
    return cell_len(plain)


def padded(value: Text, width: int) -> Text:
    return combine(value, " " * max(0, width - visible_len(value)))


def left_padded(value: Text, width: int) -> Text:
    return combine(" " * max(0, width - visible_len(value)), value)


def truncate_text(val: Text, max_w: int) -> Text:
    """Truncate a Text object while preserving Rich styles on remaining segments."""
    if visible_len(val) <= max_w:
        return val
    if max_w <= 1:
        return text("…", DIM)

    res = Text()
    curr_w = 0
    ellipsis_w = cell_len("…")
    budget = max_w - ellipsis_w

    for content, style in val.segments:
        plain_content = re.sub(r"\x1b\[[0-9;]*m", "", content)
        seg_w = cell_len(plain_content)
        if curr_w + seg_w <= budget:
            res.segments.append((plain_content, style))
            curr_w += seg_w
        else:
            avail = budget - curr_w
            truncated_chars: list[str] = []
            char_w = 0
            for char in plain_content:
                cw = cell_len(char)
                if char_w + cw <= avail:
                    truncated_chars.append(char)
                    char_w += cw
                else:
                    break
            if truncated_chars:
                res.segments.append(("".join(truncated_chars), style))
            break

    res.segments.append(("…", DIM))
    return res


def sanitize_text(
    value: str,
    *,
    strip_spaces: bool = True,
    max_chars: int | None = None,
) -> str:
    """Bound code points and remove controls before terminal rendering."""
    val_str = value[:max_chars] if max_chars is not None else value
    val_str = val_str.replace("\0", "")
    # Remove OSC sequences: \x1b]...\x07 or \x1b]...\x1b\
    val_str = re.sub(r"\x1b\][^\x07\x1b]*(\x07|\x1b\\)", "", val_str)
    # Remove SGR and other CSI escape sequences: \x1b[...]
    val_str = re.sub(r"\x1b\[[0-9;?]*[a-zA-Z]", "", val_str)
    # Unicode directional controls can visually reorder task text even though
    # they occupy no cells, so strip them with the terminal control ranges.
    val_str = re.sub(
        r"[\x00-\x1f\x7f-\x9f\u061c\u200e-\u200f\u202a-\u202e\u2066-\u2069]",
        "",
        val_str,
    )
    return val_str.strip() if strip_spaces else val_str


def json_text(
    value: JSONValue,
    *,
    strip_spaces: bool = True,
    max_chars: int | None = None,
) -> str:
    """Return sanitized scalar text; containers and booleans degrade to empty."""
    if isinstance(value, str):
        return sanitize_text(
            value, strip_spaces=strip_spaces, max_chars=max_chars
        )
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return sanitize_text(
            str(value), strip_spaces=strip_spaces, max_chars=max_chars
        )
    return ""


def parse_json_integer(value: str) -> int:
    """Parse small JSON integers and saturate oversized literals safely."""
    negative = value.startswith("-")
    digit_count = len(value) - int(negative)
    if digit_count > 18:
        return -1 if negative else MAX_INTEGER + 1
    return int(value)


def contains_surrogate(value: JSONValue) -> bool:
    """Detect lone JSON surrogates that jq rejects but json.loads accepts."""
    stack: list[JSONValue] = [value]
    while stack:
        current = stack.pop()
        if isinstance(current, str):
            if any(0xD800 <= ord(character) <= 0xDFFF for character in current):
                return True
        elif isinstance(current, dict):
            for key, item in current.items():
                if any(0xD800 <= ord(character) <= 0xDFFF for character in key):
                    return True
                stack.append(item)
        elif isinstance(current, list):
            stack.extend(current)
    return False


def number(value: object, default: Decimal = ZERO_DECIMAL) -> Decimal:
    try:
        num = Decimal(str(value))
        if num.is_nan() or num.is_infinite():
            return default
        return num
    except (InvalidOperation, ValueError, TypeError):
        return default


def integer(
    value: object,
    default: int = 0,
    *,
    maximum: int = MAX_INTEGER,
) -> int:
    """Return a bounded non-negative integer without materializing huge values."""
    amount = number(value, default=Decimal(default))
    if amount < 0:
        return max(0, default)
    if amount > maximum:
        return maximum
    return int(amount)


def format_tokens(value: object) -> str:
    """Format a token count with SI boundaries."""
    amount = number(value)
    if amount < 0:
        amount = Decimal(0)
    if amount > MAX_TOKEN_DISPLAY:
        return "999M+"
    try:
        if amount >= 1_000_000:
            millions = (amount / Decimal(1_000_000)).quantize(
                Decimal("0.1"), rounding=ROUND_HALF_UP
            )
            rendered = f"{millions.normalize():f}M"
        elif amount >= 1_000:
            thousands = (amount / Decimal(1_000)).quantize(
                Decimal("1"), rounding=ROUND_HALF_UP
            )
            if thousands >= 1_000:
                rendered = str(humanize.metric(1_000_000, precision=0))
            else:
                rendered = str(
                    humanize.metric(int(thousands * 1_000), precision=0)
                )
        else:
            rendered = str(integer(amount))
    except (InvalidOperation, ValueError, OverflowError):
        rendered = "0"
    compact = rendered.replace(" ", "")
    return compact


def join_cells(cells: list[Text]) -> Text:
    """Join non-empty cells with ' | ' separators."""
    separator = combine(" ", text("|", DIM), " ")
    result = Text()
    first = True
    for cell in cells:
        if not cell:
            continue
        if not first:
            result.append_text(separator)
        result.append_text(cell)
        first = False
    return result


def format_elapsed(start_time_ms: int) -> str:
    """Return a compact HH:MM or MM:SS elapsed string from an epoch-ms timestamp."""
    elapsed_s = max(0, int(time.time() - start_time_ms / 1000.0))
    hours, rem = divmod(elapsed_s, 3600)
    minutes, seconds = divmod(rem, 60)
    if hours > 0:
        return f"{hours}:{minutes:02d}h"
    return f"{minutes}:{seconds:02d}"


@dataclass(frozen=True)
class TaskCells:
    model: Text
    effort: Text
    elapsed: Text
    tokens: Text
    task: Text


@dataclass(frozen=True)
class ColumnWidths:
    model: int
    effort: int
    elapsed: int
    tokens: int


@dataclass(frozen=True)
class PreparedTask:
    task_id: str
    cells: TaskCells


def task_name(task: JSONObject) -> str:
    """Return the first valid task name, degrading malformed fields to ``agent``."""
    raw_name = task.get("name")
    if isinstance(raw_name, dict):
        name = json_text(raw_name.get("name"), max_chars=TASK_INPUT_CAP)
    else:
        name = json_text(raw_name, max_chars=TASK_INPUT_CAP)
    if name:
        return name

    for key in ("label", "type"):
        fallback = json_text(task.get(key), max_chars=TASK_INPUT_CAP)
        if fallback:
            return fallback
    return "agent"


def start_time_ms(value: JSONValue) -> int | None:
    """Return a plausible epoch-millisecond timestamp, or ``None`` if invalid."""
    amount = number(value)
    if not MIN_EPOCH_MILLISECONDS <= amount <= MAX_EPOCH_MILLISECONDS:
        return None
    return int(amount)


def build_task_cells(task: JSONObject) -> TaskCells:
    """Build bounded styled cells without applying batch-dependent alignment."""

    # ── task / status ─────────────────────────────────────────────────────
    name_str = task_name(task)

    status_raw = json_text(task.get("status"), max_chars=TASK_INPUT_CAP)
    if status_raw:
        cell_task = combine(
            text(name_str, WHITE), text("/", DIM), text(status_raw, DIM)
        )
    else:
        cell_task = text(name_str, WHITE)
    cell_task = truncate_text(cell_task, TASK_WIDTH_CAP)

    # ── model effort ──────────────────────────────────────────────────────
    model_value = task.get("model")
    if isinstance(model_value, dict):
        raw_model_str = json_text(
            model_value.get("display_name"), max_chars=MODEL_INPUT_CAP
        )
    else:
        raw_model_str = json_text(model_value, max_chars=MODEL_INPUT_CAP)

    model_name = re.sub(
        r"\s*\(([0-9]+(?:\.[0-9]+)?[kKmM])\s+context\)",
        r" \1",
        raw_model_str,
        count=1,
    ).strip()
    model_name = re.sub(r"^claude-", "", model_name, count=1, flags=re.IGNORECASE)
    if not model_name:
        model_name = "Claude"
    cell_model = truncate_text(text(model_name, BLUE), MODEL_WIDTH_CAP)

    effort_value = task.get("effort")
    if isinstance(effort_value, dict):
        effort_str = json_text(
            effort_value.get("level"), max_chars=EFFORT_INPUT_CAP
        )
    else:
        effort_str = json_text(effort_value, max_chars=EFFORT_INPUT_CAP)
    if effort_str:
        effort_level = effort_str.lower()
        effort_label = "med" if effort_level == "medium" else effort_level
        effort_style = {
            "low": DIM,
            "medium": ORANGE,
            "high": GREEN,
            "xhigh": PURPLE,
            "max": RED,
        }.get(effort_level, GREEN)
        cell_effort = text(effort_label, effort_style)
    else:
        cell_effort = text("-", DIM)

    thinking_value = task.get("thinking")
    thinking_on = thinking_value is True or (
        isinstance(thinking_value, dict) and thinking_value.get("enabled") is True
    )
    if thinking_on:
        cell_effort = combine(text("✦", PURPLE), " ", cell_effort)
    cell_effort = truncate_text(cell_effort, EFFORT_WIDTH_CAP)

    # ── tokens / context ──────────────────────────────────────────────────
    token_count = integer(task.get("tokenCount") or task.get("tokens") or 0)
    context_size = integer(
        task.get("contextWindowSize") or task.get("context_window_size") or 0
    )

    cell_tokens = Text()
    if token_count > 0 or context_size > 0:
        t_str = format_tokens(token_count) if token_count > 0 else "0"
        cell_tokens = truncate_text(text(t_str, ORANGE), TOKEN_WIDTH_CAP)

    # ── elapsed ───────────────────────────────────────────────────────────
    cell_elapsed = Text()
    start_ms = start_time_ms(task.get("startTime"))
    if start_ms is not None:
        cell_elapsed = truncate_text(
            text(format_elapsed(start_ms), DIM), ELAPSED_WIDTH_CAP
        )

    return TaskCells(
        model=cell_model,
        effort=cell_effort,
        elapsed=cell_elapsed,
        tokens=cell_tokens,
        task=cell_task,
    )


def measure_columns(tasks: list[TaskCells]) -> ColumnWidths:
    """Return display-cell widths shared by every row in a task batch."""
    return ColumnWidths(
        model=min(
            max((visible_len(task.model) for task in tasks), default=0),
            MODEL_WIDTH_CAP,
        ),
        effort=min(
            max((visible_len(task.effort) for task in tasks), default=0),
            EFFORT_WIDTH_CAP,
        ),
        elapsed=min(
            max((visible_len(task.elapsed) for task in tasks), default=0),
            ELAPSED_WIDTH_CAP,
        ),
        tokens=min(
            max((visible_len(task.tokens) for task in tasks), default=0),
            TOKEN_WIDTH_CAP,
        ),
    )


def prepare_task(value: JSONValue) -> PreparedTask | None:
    """Validate one task and return bounded cells without affecting siblings.

    Non-object tasks, lone surrogates, and missing identifiers degrade to no
    override. All other malformed fields degrade independently inside their
    cells, so one unresolved task cannot suppress valid rows in the batch.
    """
    if not isinstance(value, dict) or contains_surrogate(value):
        return None
    task_id = json_text(value.get("id"), strip_spaces=False)
    if not task_id:
        return None
    return PreparedTask(task_id=task_id, cells=build_task_cells(value))


def render_task(task_cells: TaskCells, widths: ColumnWidths, max_columns: int) -> str:
    """Render one row using widths measured across its complete input batch."""
    cell_model_effort = combine(
        padded(task_cells.model, widths.model),
        " ",
        left_padded(task_cells.effort, widths.effort),
    )
    row_cells: list[Text] = [
        cell_model_effort,
        left_padded(task_cells.elapsed, widths.elapsed),
        left_padded(task_cells.tokens, widths.tokens),
        task_cells.task,
    ]
    line = join_cells(row_cells)

    if max_columns > 0:
        candidates = [
            row_cells,
            [cell_model_effort, row_cells[2], row_cells[3]],
            [cell_model_effort, row_cells[3]],
        ]

        # Every row chooses its tier from the same aligned prefix width. Choosing
        # only after a particular row overflows would let short tasks keep pipes
        # that longer tasks dropped.
        for candidate_cells in candidates:
            prefix_cells = candidate_cells[:-1]
            prefix_line = join_cells(prefix_cells)
            separator_len = 3 if prefix_line and row_cells[3] else 0
            avail_for_task = max_columns - visible_len(prefix_line) - separator_len

            if avail_for_task >= 6 or candidate_cells == candidates[-1]:
                target_w = max(6, avail_for_task)
                truncated_task = truncate_text(row_cells[3], target_w)
                line = join_cells(prefix_cells + [truncated_task])
                break

        if visible_len(line) > max_columns:
            line = truncate_text(line, max_columns)

    return line.render()


def silence_stdout() -> None:
    """Prevent interpreter shutdown from retrying a broken stdout pipe."""
    try:
        devnull_fd = os.open(os.devnull, os.O_WRONLY)
        try:
            os.dup2(devnull_fd, sys.stdout.fileno())
        finally:
            os.close(devnull_fd)
    except OSError:
        pass


def main() -> int:
    """Read stdin and emit JSON output line(s) for subagent tasks.

    Malformed documents and parser recursion limits degrade to no overrides
    with a successful exit. Raising the interpreter recursion limit is unsafe:
    it can exchange a clean empty result for a process-level C-stack failure.
    """
    input_text = (
        sys.stdin.buffer.read().decode("utf-8", errors="surrogateescape").rstrip("\n")
    )
    if not input_text:
        return 0

    try:
        parsed = cast(
            JSONValue,
            json.loads(input_text, parse_int=parse_json_integer),
        )
    except (ValueError, RecursionError):
        return 0

    max_columns = integer(os.environ.get("COLUMNS") or 0, maximum=MAX_COLUMNS)
    raw_tasks: list[JSONValue] = []
    if isinstance(parsed, dict):
        max_columns = integer(
            parsed.get("columns"),
            default=max_columns,
            maximum=MAX_COLUMNS,
        )
        t_val = parsed.get("tasks")
        if isinstance(t_val, list):
            raw_tasks = t_val
        elif "id" in parsed or "model" in parsed or "cwd" in parsed:
            raw_tasks = [parsed]
    elif isinstance(parsed, list):
        raw_tasks = parsed

    prepared_tasks: list[PreparedTask] = []
    for raw_task in raw_tasks:
        prepared = prepare_task(raw_task)
        if prepared is not None:
            prepared_tasks.append(prepared)
    widths = measure_columns([task.cells for task in prepared_tasks])

    for task in prepared_tasks:
        content = render_task(task.cells, widths, max_columns)
        out_json = json.dumps(
            {"id": task.task_id, "content": content}, separators=(",", ":")
        )
        try:
            _ = sys.stdout.buffer.write(
                (out_json + "\n").encode("utf-8", errors="surrogateescape")
            )
        except BrokenPipeError:
            silence_stdout()
            return 0

    try:
        sys.stdout.buffer.flush()
    except BrokenPipeError:
        silence_stdout()

    return 0


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