tau-coding-agent 0.1.0__py3-none-any.whl

tau/tui/agent_hooks.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from tau.runtime.service import Runtime
+    from tau.tui.components.layout import Layout
+    from tau.tui.components.message_list import MessageBlock
+    from tau.tui.tui import TUI
+
+# Flush streamed tokens to the block at most once per frame (~60fps).
+# Markdown is re-parsed only on each flush, not on every token.
+_STREAM_FLUSH_INTERVAL = 1 / 60
+
+
+def _find_component(root: object, attr: str) -> object | None:
+    """Depth-first search the component tree for one exposing ``attr``."""
+    if root is None:
+        return None
+    if hasattr(root, attr):
+        return root
+    for child in (getattr(root, "children", None) or []):
+        found = _find_component(child, attr)
+        if found is not None:
+            return found
+    for slot in (getattr(root, "_slots", None) or []):
+        comp = slot[0] if isinstance(slot, tuple) else slot
+        found = _find_component(comp, attr)
+        if found is not None:
+            return found
+    return None
+
+
+class AgentHookHandler:
+    """Subscribes to agent lifecycle hooks and drives the TUI in response.
+
+    Owns all streaming state (current block, terminal block) so App does not
+    need to track it. Call ``subscribe()`` after the agent is ready, then
+    ``unsubscribe()`` (or use the returned unsub callables) on teardown.
+    """
+
+    def __init__(self, runtime: Runtime, layout: Layout, tui: TUI,
+                 on_palette_refresh: Callable[[], None] | None = None,
+                 on_turn_content: Callable[[], None] | None = None) -> None:
+        self._runtime = runtime
+        self._layout = layout
+        self._tui = tui
+        self._on_palette_refresh = on_palette_refresh
+        self._on_turn_content = on_turn_content
+
+        self._current_block: MessageBlock | None = None
+        self._current_text_length: int = 0
+        self._current_terminal_block: MessageBlock | None = None
+        self._unsubs: list[Callable[[], None]] = []
+
+        # Streaming batch state — pending token flush
+        self._pending_msg: object = None
+        self._pending_flush_handle: asyncio.TimerHandle | None = None
+        self._last_flush_at: float = 0.0
+
+    def subscribe(self) -> None:
+        """Register all hook handlers on the current agent."""
+        agent = self._runtime.agent
+        if agent is None:
+            return
+        hooks = agent.hooks
+        self._unsubs = [
+            hooks.register("agent_start",          self._on_agent_start),
+            hooks.register("agent_end",            self._on_agent_end),
+            hooks.register("settled",              self._on_settled),
+            hooks.register("message_start",        self._on_message_start),
+            hooks.register("message_update",       self._on_message_update),
+            hooks.register("message_end",          self._on_message_end),
+            hooks.register("message_rollback",     self._on_message_rollback),
+            hooks.register("tool_execution_start", self._on_tool_start),
+            hooks.register("tool_execution_end",   self._on_tool_end),
+            hooks.register("model_select",         self._on_model_select),
+            hooks.register("terminal_execution",       self._on_terminal_execution),
+            hooks.register("terminal_output",          self._on_terminal_output),
+            hooks.register("session_start",        self._on_session_start),
+            hooks.register("queue_update",         self._on_queue_update),
+            hooks.register("compaction_start",     self._on_compaction_start),
+            hooks.register("compaction_end",       self._on_compaction_end),
+        ]
+
+    def unsubscribe(self) -> None:
+        for unsub in self._unsubs:
+            unsub()
+        self._unsubs.clear()
+
+    # ── Session ───────────────────────────────────────────────────────────────
+
+    async def _on_session_start(self, event: object) -> None:
+        from tau.hooks.session import SessionStartReason
+        reason = getattr(event, "reason", None)
+        # Reasons that swap in a different message history: Fork/Resume/Clone
+        # replay the new branch into the transcript; New just clears it.
+        replay = reason in (
+            SessionStartReason.Fork,
+            SessionStartReason.Resume,
+            SessionStartReason.Clone,
+        )
+        if reason == SessionStartReason.New or replay:
+            self._layout.clear_messages()
+            self._layout.spinner.stop()
+            self._current_block = None
+            self._current_text_length = 0
+            self._current_terminal_block = None
+            if replay:
+                sm = self._runtime.session_manager
+                if sm is not None:
+                    ctx = sm.build_session_context()
+                    for msg in ctx.messages:
+                        self._layout.add_message(msg)
+        sm = self._runtime.session_manager
+        if sm is not None:
+            self._layout.set_cwd(sm.cwd)
+        # Resume/branch swap in a different message history but run no agent turn,
+        # so the footer's context-usage badge would otherwise keep the previous
+        # session's value. Push fresh usage to it from this guaranteed refresh point.
+        self._refresh_footer_context()
+        if self._on_palette_refresh is not None:
+            self._on_palette_refresh()
+        self._tui.request_render()
+
+    def _refresh_footer_context(self) -> None:
+        """Re-push current context usage to the footer model badge, if present."""
+        footer = getattr(self._layout, "footer", None)
+        badge = _find_component(footer, "update_context_from_ctx")
+        update = getattr(badge, "update_context_from_ctx", None)
+        if not callable(update):
+            return
+        try:
+            from tau.extensions.context import ExtensionContext
+            update(ExtensionContext.from_runtime(self._runtime))
+        except Exception:
+            pass
+
+    # ── Agent lifecycle ───────────────────────────────────────────────────────
+
+    async def _on_agent_start(self, _event: object) -> None:
+        self._spinner(self._layout.spinner._theme.label_thinking, running=True)
+
+    async def _on_agent_end(self, _event: object) -> None:
+        self._spinner(running=False)
+
+    async def _on_settled(self, _event: object) -> None:
+        self._layout.set_pending_queue([], [])
+        self._spinner(running=False)
+
+    async def _on_compaction_start(self, _event: object) -> None:
+        self._spinner(self._layout.spinner._theme.label_compacting, running=True)
+
+    async def _on_compaction_end(self, _event: object) -> None:
+        self._spinner(running=False)
+
+    # ── Messages ──────────────────────────────────────────────────────────────
+
+    async def _on_message_start(self, event: object) -> None:
+        msg = getattr(event, "message", None)
+        if msg is None:
+            return
+        self._layout.spinner.set_label(self._layout.spinner._theme.label_thinking)
+        block = self._layout.add_message(msg, streaming=False)
+        self._current_block = block
+        self._current_text_length = _text_length(msg)
+        self._tui.request_render()
+
+    async def _on_message_update(self, event: object) -> None:
+        from tau.message.types import ThinkingContent, TextContent
+        msg = getattr(event, "message", None)
+        if msg is None or self._current_block is None:
+            return
+
+        # Update spinner label immediately — zero cost.
+        contents = getattr(msg, "contents", [])
+        if contents:
+            self._mark_turn_content()
+            last = contents[-1]
+            if isinstance(last, ThinkingContent):
+                self._layout.spinner.set_label(self._layout.spinner._theme.label_thinking)
+            elif isinstance(last, TextContent) and last.content:
+                self._layout.spinner.set_label(self._layout.spinner._theme.label_streaming)
+
+        # Buffer the latest message; schedule a flush if none pending.
+        self._pending_msg = msg
+        if self._pending_flush_handle is None:
+            elapsed = time.monotonic() - self._last_flush_at
+            delay = max(0.0, _STREAM_FLUSH_INTERVAL - elapsed)
+            loop = asyncio.get_event_loop()
+            self._pending_flush_handle = loop.call_later(delay, self._flush_pending)
+
+    def _flush_pending(self) -> None:
+        """Flush the buffered token batch: re-parse markdown once, then render."""
+        self._pending_flush_handle = None
+        msg = self._pending_msg
+        self._pending_msg = None
+        if msg is None or self._current_block is None:
+            return
+        tl = _text_length(msg)
+        self._update_block(msg, streaming=tl > self._current_text_length)
+        self._current_text_length = tl
+        self._last_flush_at = time.monotonic()
+        self._tui.request_render()
+
+    async def _on_message_end(self, event: object) -> None:
+        from tau.message.types import AssistantMessage, ToolMessage
+        # Cancel any pending batch flush — message_end supersedes it.
+        if self._pending_flush_handle is not None:
+            self._pending_flush_handle.cancel()
+            self._pending_flush_handle = None
+        self._pending_msg = None
+        self._last_flush_at = 0.0
+
+        msg = getattr(event, "message", None)
+        if msg is None:
+            return
+        if isinstance(msg, (AssistantMessage, ToolMessage)):
+            self._mark_turn_content()
+            if self._current_block is not None:
+                self._update_block(msg, streaming=False, clear=True)
+            else:
+                self._layout.add_message(msg)
+        self._tui.request_render()
+
+    async def _on_message_rollback(self, event: object) -> None:
+        """Remove the last ``count`` message blocks from the transcript.
+
+        Mirrors the engine dropping an interrupted tool turn: cancel any pending
+        flush, drop the assistant tool-call block and its tool-result block, and
+        reset the streaming cursor so the following interrupt marker renders fresh.
+        """
+        count = getattr(event, "count", 0)
+        if self._pending_flush_handle is not None:
+            self._pending_flush_handle.cancel()
+            self._pending_flush_handle = None
+        self._pending_msg = None
+        for _ in range(count):
+            if not self._layout.messages.remove_last():
+                break
+        self._current_block = None
+        self._current_text_length = 0
+        self._tui.request_render()
+
+    # ── Tools ─────────────────────────────────────────────────────────────────
+
+    async def _on_tool_start(self, _event: object) -> None:
+        self._mark_turn_content()
+        self._spinner(self._layout.spinner._theme.label_tool_calling)
+
+    async def _on_tool_end(self, _event: object) -> None:
+        self._spinner(self._layout.spinner._theme.label_thinking)
+
+    # ── Terminal ──────────────────────────────────────────────────────────
+
+    async def _on_terminal_execution(self, event: object) -> None:
+        msg = getattr(event, "message", None)
+        streaming = getattr(event, "streaming", False)
+        if msg is None:
+            return
+        if streaming:
+            block = self._layout.add_message(msg, streaming=True)
+            self._current_terminal_block = block
+        else:
+            if self._current_terminal_block is not None:
+                self._current_terminal_block.set_streaming(False)
+                self._current_terminal_block = None
+        self._tui.request_render()
+
+    async def _on_terminal_output(self, _event: object) -> None:
+        if self._current_terminal_block is not None:
+            self._current_terminal_block.invalidate()
+            self._tui.request_render()
+
+    # ── Model / queue ─────────────────────────────────────────────────────────
+
+    async def _on_model_select(self, _event: object) -> None:
+        if self._on_palette_refresh is not None:
+            self._on_palette_refresh()
+        self._tui.request_render()
+
+    async def _on_queue_update(self, event: object) -> None:
+        from tau.message.types import TextContent
+
+        engine_state = None
+        agent = self._runtime.agent
+        if agent is not None:
+            engine_state = agent._engine.state
+
+        steering: list[str] = []
+        followup:  list[str] = []
+
+        if engine_state is not None:
+            for queue, out in [
+                (engine_state.steering_queue, steering),
+                (engine_state.follow_up_queue, followup),
+            ]:
+                if queue:
+                    for msg in queue.snapshot():
+                        text = "".join(
+                            c.content for c in getattr(msg, "contents", [])
+                            if isinstance(c, TextContent)
+                        )
+                        if text:
+                            out.append(text)
+
+        self._layout.set_pending_queue(steering, followup)
+        self._tui.request_render()
+
+    # ── Helpers ───────────────────────────────────────────────────────────────
+
+    def _mark_turn_content(self) -> None:
+        """Tell the input handler the assistant has produced output this turn."""
+        if self._on_turn_content is not None:
+            self._on_turn_content()
+
+    def _spinner(self, label: str | None = None, *, running: bool | None = None) -> None:
+        if label is not None:
+            self._layout.spinner.set_label(label)
+        if running is True:
+            self._layout.spinner.start()
+        elif running is False:
+            self._layout.spinner.stop()
+        self._tui.request_render()
+
+    def _update_block(self, msg: object, *, streaming: bool, clear: bool = False) -> None:
+        if self._current_block is None:
+            return
+        self._current_block._message = msg
+        self._current_block.set_streaming(streaming)
+        self._current_block.invalidate()
+        if clear:
+            self._current_block = None
+            self._current_text_length = 0
+
+    def _refresh_model_badge(self) -> None:
+        if self._on_palette_refresh is not None:
+            self._on_palette_refresh()
+
+
+def _text_length(message: object) -> int:
+    from tau.message.types import TextContent
+    contents = getattr(message, "contents", [])
+    return sum(len(item.content) for item in contents if isinstance(item, TextContent))

tau/tui/ansi.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+import re
+import unicodedata
+
+# ── Regex to strip all ANSI escape sequences ─────────────────────────────────
+_ANSI_RE = re.compile(
+    r"\x1b(?:"
+    r"\[[0-9;:<=>?]*[ -/]*[@-~]"   # CSI sequences
+    r"|\][^\x07\x1b]*(?:\x07|\x1b\\)"  # OSC sequences
+    r"|_[^\x1b]*(?:\x1b\\|\x07)"   # APC sequences
+    r"|[PX^][^\x1b]*\x1b\\"        # DCS / PM / SOS
+    r"|[@-_]"                       # 2-char Fe sequences (ESC @-_)
+    r")"
+)
+
+# ── Cursor marker ────────────────────────────────────────────────────────────
+# An APC sequence (zero visible width, stripped by _ANSI_RE) injected into a
+# rendered line to tell the Renderer where to position the hardware cursor for
+# IME candidate windows.  TextInput inserts it; Renderer removes it and moves
+# the real terminal cursor to the marked column.
+CURSOR_MARKER = "\x1b_C\x1b\\"
+
+# ── SGR (colour / style) ──────────────────────────────────────────────────────
+
+RESET      = "\x1b[0m"
+BOLD       = "\x1b[1m"
+DIM        = "\x1b[2m"
+ITALIC     = "\x1b[3m"
+UNDERLINE  = "\x1b[4m"
+BLINK      = "\x1b[5m"
+REVERSE    = "\x1b[7m"
+STRIKE     = "\x1b[9m"
+
+# Standard foreground colours
+BLACK   = "\x1b[30m"
+RED     = "\x1b[31m"
+GREEN   = "\x1b[32m"
+YELLOW  = "\x1b[33m"
+BLUE    = "\x1b[34m"
+MAGENTA = "\x1b[35m"
+CYAN    = "\x1b[36m"
+WHITE   = "\x1b[37m"
+DEFAULT = "\x1b[39m"
+
+# Bright foreground colours
+BRIGHT_BLACK   = "\x1b[90m"
+BRIGHT_RED     = "\x1b[91m"
+BRIGHT_GREEN   = "\x1b[92m"
+BRIGHT_YELLOW  = "\x1b[93m"
+BRIGHT_BLUE    = "\x1b[94m"
+BRIGHT_MAGENTA = "\x1b[95m"
+BRIGHT_CYAN    = "\x1b[96m"
+BRIGHT_WHITE   = "\x1b[97m"
+
+# Standard background colours
+BG_BLACK   = "\x1b[40m"
+BG_RED     = "\x1b[41m"
+BG_GREEN   = "\x1b[42m"
+BG_YELLOW  = "\x1b[43m"
+BG_BLUE    = "\x1b[44m"
+BG_MAGENTA = "\x1b[45m"
+BG_CYAN    = "\x1b[46m"
+BG_WHITE   = "\x1b[47m"
+BG_DEFAULT = "\x1b[49m"
+
+
+def fg(r: int, g: int, b: int) -> str:
+    """Truecolor foreground."""
+    return f"\x1b[38;2;{r};{g};{b}m"
+
+
+def bg(r: int, g: int, b: int) -> str:
+    """Truecolor background."""
+    return f"\x1b[48;2;{r};{g};{b}m"
+
+
+def fg256(code: int) -> str:
+    """256-colour foreground (0–255)."""
+    return f"\x1b[38;5;{code}m"
+
+
+def bg256(code: int) -> str:
+    """256-colour background (0–255)."""
+    return f"\x1b[48;5;{code}m"
+
+
+def style(text: str, *codes: str) -> str:
+    """Wrap text with one or more SGR codes, reset at end."""
+    return "".join(codes) + text + RESET
+
+
+# ── Width calculation ─────────────────────────────────────────────────────────
+
+def _char_width(ch: str) -> int:
+    """Return the terminal column width of a single character."""
+    cp = ord(ch)
+    # Null byte / C0 / C1 control chars
+    if cp == 0 or (0x007F <= cp <= 0x009F):
+        return 0
+    eaw = unicodedata.east_asian_width(ch)
+    if eaw in ("W", "F"):
+        return 2
+    if eaw == "Na" or unicodedata.category(ch) in ("Mn", "Me", "Cf"):
+        # Combining marks and format chars have zero width
+        if unicodedata.category(ch) in ("Mn", "Me", "Cf"):
+            return 0
+    return 1
+
+
+def visible_width(text: str) -> int:
+    """Return the number of terminal columns the string will occupy."""
+    # Fast path: pure ASCII printable
+    if text.isascii() and text.isprintable():
+        return len(text)
+    stripped = strip_ansi(text)
+    return sum(_char_width(ch) for ch in stripped)
+
+
+# ── Strip ─────────────────────────────────────────────────────────────────────
+
+def strip_ansi(text: str) -> str:
+    """Remove all ANSI escape sequences from text."""
+    return _ANSI_RE.sub("", text)
+
+
+# ── Truncation ────────────────────────────────────────────────────────────────
+
+def truncate(text: str, max_width: int, ellipsis: str = "…") -> str:
+    """
+    Truncate text to max_width columns, preserving ANSI codes.
+    Appends ellipsis if truncation occurred.
+    """
+    if visible_width(text) <= max_width:
+        return text
+
+    ellipsis_w = visible_width(ellipsis)
+    target = max_width - ellipsis_w
+    result, _ = _take_columns(text, target)
+    return result + RESET + ellipsis + RESET
+
+
+def pad(text: str, width: int, char: str = " ", align: str = "left") -> str:
+    """
+    Pad text to exactly width columns with char.
+    align: 'left' | 'right' | 'center'
+    """
+    current = visible_width(text)
+    deficit = max(0, width - current)
+    if align == "right":
+        return char * deficit + text
+    if align == "center":
+        left = deficit // 2
+        right = deficit - left
+        return char * left + text + char * right
+    return text + char * deficit
+
+
+# ── Wrapping ──────────────────────────────────────────────────────────────────
+
+def wrap(text: str, width: int) -> list[str]:
+    """
+    Wrap text to width columns, preserving ANSI escape codes across line breaks.
+    Hard newlines in the source are respected.
+    """
+    if width <= 0:
+        return [text]
+
+    result: list[str] = []
+    for source_line in text.split("\n"):
+        result.extend(_wrap_single_line(source_line, width))
+    return result
+
+
+def _wrap_single_line(line: str, width: int) -> list[str]:
+    """Wrap a single line (no embedded newlines) to width columns."""
+    if visible_width(line) <= width:
+        return [line]
+
+    lines: list[str] = []
+    tracker = _AnsiStateTracker()
+    remaining = line
+
+    while remaining:
+        chunk, remaining = _split_at_columns(remaining, width, tracker)
+        lines.append(chunk)
+
+    return lines
+
+
+def _split_at_columns(text: str, width: int, tracker: _AnsiStateTracker) -> tuple[str, str]:
+    """
+    Split text into (head, tail) where head fits within width columns.
+    Tracker maintains active SGR state so head starts with the right codes
+    and tail can be resumed correctly.
+    """
+    prefix = tracker.active_codes()
+    taken, col = "", 0
+    i = 0
+
+    while i < len(text):
+        # Check for ANSI escape at current position
+        m = _ANSI_RE.match(text, i)
+        if m:
+            code = m.group(0)
+            tracker.process(code)
+            taken += code
+            i += len(code)
+            continue
+
+        ch = text[i]
+        w = _char_width(ch)
+        if col + w > width:
+            break
+        taken += ch
+        col += w
+        i += 1
+
+    head = prefix + taken + (RESET if tracker.has_state() else "")
+    tail = text[i:]
+    return head, tail
+
+
+def _take_columns(text: str, max_cols: int) -> tuple[str, int]:
+    """Return (text_up_to_max_cols, actual_cols_taken), preserving ANSI codes."""
+    result, col = "", 0
+    i = 0
+    while i < len(text):
+        m = _ANSI_RE.match(text, i)
+        if m:
+            result += m.group(0)
+            i += len(m.group(0))
+            continue
+        ch = text[i]
+        w = _char_width(ch)
+        if col + w > max_cols:
+            break
+        result += ch
+        col += w
+        i += 1
+    return result, col
+
+
+# ── ANSI state tracker ────────────────────────────────────────────────────────
+
+class _AnsiStateTracker:
+    """
+    Tracks active SGR (colour/style) state so it can be re-applied after
+    a line break during wrapping.
+    """
+
+    def __init__(self) -> None:
+        self._bold = False
+        self._dim = False
+        self._italic = False
+        self._underline = False
+        self._fg: str | None = None
+        self._bg: str | None = None
+
+    def process(self, code: str) -> None:
+        m = re.match(r"\x1b\[([\d;]*)m$", code)
+        if not m:
+            return
+        params = m.group(1)
+        if not params or params == "0":
+            self._reset()
+            return
+        nums = [int(x) for x in params.split(";") if x]
+        i = 0
+        while i < len(nums):
+            n = nums[i]
+            if n == 0:
+                self._reset()
+            elif n == 1:
+                self._bold = True
+            elif n == 2:
+                self._dim = True
+            elif n == 3:
+                self._italic = True
+            elif n == 4:
+                self._underline = True
+            elif n == 22:
+                self._bold = self._dim = False
+            elif n == 23:
+                self._italic = False
+            elif n == 24:
+                self._underline = False
+            elif 30 <= n <= 37 or n == 39 or 90 <= n <= 97:
+                self._fg = f"\x1b[{n}m"
+            elif n == 38 and i + 2 < len(nums) and nums[i + 1] == 5:
+                self._fg = f"\x1b[38;5;{nums[i + 2]}m"
+                i += 2
+            elif n == 38 and i + 4 < len(nums) and nums[i + 1] == 2:
+                self._fg = f"\x1b[38;2;{nums[i+2]};{nums[i+3]};{nums[i+4]}m"
+                i += 4
+            elif 40 <= n <= 47 or n == 49 or 100 <= n <= 107:
+                self._bg = f"\x1b[{n}m"
+            elif n == 48 and i + 2 < len(nums) and nums[i + 1] == 5:
+                self._bg = f"\x1b[48;5;{nums[i + 2]}m"
+                i += 2
+            elif n == 48 and i + 4 < len(nums) and nums[i + 1] == 2:
+                self._bg = f"\x1b[48;2;{nums[i+2]};{nums[i+3]};{nums[i+4]}m"
+                i += 4
+            i += 1
+
+    def _reset(self) -> None:
+        self._bold = self._dim = self._italic = self._underline = False
+        self._fg = self._bg = None
+
+    def has_state(self) -> bool:
+        return bool(self._bold or self._dim or self._italic or
+                    self._underline or self._fg or self._bg)
+
+    def active_codes(self) -> str:
+        if not self.has_state():
+            return ""
+        parts = []
+        if self._bold:
+            parts.append(BOLD)
+        if self._dim:
+            parts.append(DIM)
+        if self._italic:
+            parts.append(ITALIC)
+        if self._underline:
+            parts.append(UNDERLINE)
+        if self._fg:
+            parts.append(self._fg)
+        if self._bg:
+            parts.append(self._bg)
+        return "".join(parts)