vibe/tests/snapshots/test_ui_snapshot_tool_hooks.py

from __future__ import annotations

from typing import cast

from textual.app import App, ComposeResult
from textual.containers import VerticalScroll
from textual.pilot import Pilot
from textual.widget import Widget

from tests.snapshots.snap_compare import SnapCompare
from tests.stubs.fake_tool import FakeTool, FakeToolArgs, FakeToolResult
from vibe.cli.textual_ui.handlers.event_handler import EventHandler
from vibe.cli.textual_ui.widgets.collapsible import CollapsibleSection
from vibe.cli.textual_ui.widgets.tools import ToolCallMessage
from vibe.core.hooks.models import (
    HookEndEvent,
    HookMessageSeverity,
    HookRunEndEvent,
    HookRunStartEvent,
    HookStartEvent,
    HookType,
)
from vibe.core.types import ToolCallEvent, ToolResultEvent


class ToolHooksApp(App):
    CSS_PATH = "../../vibe/cli/textual_ui/app.tcss"

    def __init__(self) -> None:
        super().__init__()
        self._scroll: VerticalScroll | None = None
        self._handler: EventHandler | None = None

    def compose(self) -> ComposeResult:
        self._scroll = VerticalScroll(id="messages")
        yield self._scroll

    def on_mount(self) -> None:
        async def mount_callback(
            widget: Widget, *, after: Widget | None = None, before: Widget | None = None
        ) -> None:
            if self._scroll is None:
                return
            if before is not None and before.parent is self._scroll:
                await self._scroll.mount(widget, before=before)
            elif after is not None and after.parent is self._scroll:
                await self._scroll.mount(widget, after=after)
            else:
                await self._scroll.mount(widget)

        self._handler = EventHandler(
            mount_callback=mount_callback, get_tools_collapsed=lambda: False
        )

    def freeze_spinners(self) -> None:
        for widget in self.query(ToolCallMessage):
            widget._is_spinning = False
            if widget._spinner_timer:
                widget._spinner_timer.stop()
                widget._spinner_timer = None
            widget._spinner.reset()
            if widget._indicator_widget:
                widget._indicator_widget.update(widget._spinner.current_frame())

    async def emit_tool_with_hooks(self) -> None:
        if self._handler is None:
            return
        call_id = "call_1"

        # before_tool hooks
        await self._handler.handle_event(
            ToolCallEvent(
                tool_call_id=call_id,
                tool_name="stub_tool",
                tool_class=FakeTool,
                args=FakeToolArgs(),
            )
        )
        await self._handler.handle_event(
            HookRunStartEvent(
                scope=HookType.BEFORE_TOOL, tool_name="stub_tool", tool_call_id=call_id
            )
        )
        await self._handler.handle_event(HookStartEvent(hook_name="guard-bash"))
        await self._handler.handle_event(
            HookEndEvent(
                hook_name="guard-bash",
                status=HookMessageSeverity.OK,
                content="Command allowed",
                scope=HookType.BEFORE_TOOL,
                tool_call_id=call_id,
            )
        )
        await self._handler.handle_event(
            HookRunEndEvent(scope=HookType.BEFORE_TOOL, tool_call_id=call_id)
        )

        # tool result
        await self._handler.handle_event(
            ToolResultEvent(
                tool_name="stub_tool",
                tool_class=FakeTool,
                result=FakeToolResult(message="fake tool executed"),
                tool_call_id=call_id,
            )
        )

        # after_tool hooks
        await self._handler.handle_event(
            HookRunStartEvent(
                scope=HookType.AFTER_TOOL, tool_name="stub_tool", tool_call_id=call_id
            )
        )
        await self._handler.handle_event(HookStartEvent(hook_name="redact-secrets"))
        await self._handler.handle_event(
            HookEndEvent(
                hook_name="redact-secrets",
                status=HookMessageSeverity.WARNING,
                content="Replaced tool result (42 chars)",
                scope=HookType.AFTER_TOOL,
                tool_call_id=call_id,
            )
        )
        await self._handler.handle_event(
            HookRunEndEvent(scope=HookType.AFTER_TOOL, tool_call_id=call_id)
        )


def test_snapshot_tool_hooks(snap_compare: SnapCompare) -> None:
    async def run_before(pilot: Pilot) -> None:
        app = cast(ToolHooksApp, pilot.app)
        await app.emit_tool_with_hooks()
        await pilot.pause(0.3)
        app.freeze_spinners()
        await pilot.pause(0.1)

    assert snap_compare(
        "test_ui_snapshot_tool_hooks.py:ToolHooksApp",
        terminal_size=(80, 20),
        run_before=run_before,
    )


def test_snapshot_tool_hooks_expanded(snap_compare: SnapCompare) -> None:
    async def run_before(pilot: Pilot) -> None:
        app = cast(ToolHooksApp, pilot.app)
        await app.emit_tool_with_hooks()
        await pilot.pause(0.3)
        app.freeze_spinners()
        for section in app.query(CollapsibleSection):
            section.set_collapsed(False)
        await pilot.pause(0.1)

    assert snap_compare(
        "test_ui_snapshot_tool_hooks.py:ToolHooksApp",
        terminal_size=(80, 20),
        run_before=run_before,
    )