vibe/vibe/cli/textual_ui/handlers/event_handler.py

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import TYPE_CHECKING

from textual.widget import Widget

from vibe.cli.textual_ui.widgets.compact import CompactMessage
from vibe.cli.textual_ui.widgets.loading import DEFAULT_LOADING_STATUS
from vibe.cli.textual_ui.widgets.messages import (
    AssistantMessage,
    HookRunContainer,
    HookSystemMessageLine,
    PlanFileMessage,
    ReasoningMessage,
)
from vibe.cli.textual_ui.widgets.no_markup_static import NoMarkupStatic
from vibe.cli.textual_ui.widgets.tools import ToolCallMessage, ToolResultMessage
from vibe.core.hooks.models import (
    HookEndEvent,
    HookEvent,
    HookRunEndEvent,
    HookRunStartEvent,
    HookStartEvent,
    HookType,
)
from vibe.core.tools.ui import ToolUIDataAdapter
from vibe.core.types import (
    AgentProfileChangedEvent,
    AssistantEvent,
    BaseEvent,
    CompactEndEvent,
    CompactStartEvent,
    PlanReviewEndedEvent,
    PlanReviewRequestedEvent,
    ReasoningEvent,
    SessionTitleUpdatedEvent,
    ToolCallEvent,
    ToolResultEvent,
    ToolStreamEvent,
    UserMessageEvent,
    WaitingForInputEvent,
)
from vibe.core.utils import TaggedText

if TYPE_CHECKING:
    from vibe.cli.textual_ui.widgets.loading import LoadingWidget


class EventHandler:
    def __init__(
        self,
        mount_callback: Callable,
        get_tools_collapsed: Callable[[], bool],
        on_profile_changed: Callable[[], None] | None = None,
    ) -> None:
        self.mount_callback = mount_callback
        self.get_tools_collapsed = get_tools_collapsed
        self.on_profile_changed = on_profile_changed
        self.tool_calls: dict[str, ToolCallMessage] = {}
        self.current_compact: CompactMessage | None = None
        self.current_streaming_message: AssistantMessage | None = None
        self.current_streaming_reasoning: ReasoningMessage | None = None
        self.plan_file_message: PlanFileMessage | None = None
        # Keyed by "agent_turn", "before_tool:{call_id}", "after_tool:{call_id}"
        self._hook_containers: dict[str, HookRunContainer] = {}
        # Per-tool-call anchor for correct widget ordering during concurrent calls.
        self._tool_call_anchors: dict[str, Widget] = {}
        # Errored results shown muted while their verdict is unknown: a follow-up
        # tool call confirms them recoverable (stay muted), turn end escalates
        # them to hard errors.
        self._pending_error_results: list[ToolResultMessage] = []

    async def _handle_hook_event(
        self, event: HookEvent, loading_widget: LoadingWidget | None = None
    ) -> None:
        match event:
            case HookRunStartEvent():
                await self._handle_hook_run_start(event)
            case HookRunEndEvent():
                await self._handle_hook_run_end(event)
            case HookStartEvent():
                await self.finalize_streaming()
                if loading_widget:
                    loading_widget.set_status(f"Running hook {event.hook_name}")
            case HookEndEvent():
                if event.content:
                    key = self._hook_container_key(event.scope, event.tool_call_id)
                    container = self._hook_containers.get(key)
                    if container is not None:
                        widget = HookSystemMessageLine(
                            hook_name=event.hook_name,
                            content=event.content,
                            severity=event.status,
                        )
                        await container.add_message(widget)
                        if event.scope == HookType.BEFORE_TOOL and event.tool_call_id:
                            tool_call = self.tool_calls.get(event.tool_call_id)
                            if tool_call is not None:
                                tool_call.add_class("no-gap")
                if loading_widget:
                    loading_widget.set_status(DEFAULT_LOADING_STATUS)

    @staticmethod
    def _hook_container_key(scope: HookType, tool_call_id: str | None) -> str:
        if scope == HookType.POST_AGENT_TURN:
            return "agent_turn"
        return f"{scope.value}:{tool_call_id or ''}"

    async def _handle_hook_run_start(self, event: HookRunStartEvent) -> None:
        container = HookRunContainer()
        key = self._hook_container_key(event.scope, event.tool_call_id)
        self._hook_containers[key] = container

        if event.scope == HookType.BEFORE_TOOL:
            container.add_class("hook-before-tool")

        anchor = self._tool_call_anchors.get(event.tool_call_id or "")
        if event.scope == HookType.BEFORE_TOOL and anchor is not None:
            # Mount *above* the tool call widget so it reads as a precondition.
            await self.mount_callback(container, before=anchor)
        elif event.scope == HookType.AFTER_TOOL and anchor is not None:
            await self.mount_callback(container, after=anchor)
        else:
            await self.mount_callback(container)

    async def _handle_hook_run_end(self, event: HookRunEndEvent) -> None:
        key = self._hook_container_key(event.scope, event.tool_call_id)
        container = self._hook_containers.pop(key, None)
        if container is None:
            return
        if container.display:
            # BEFORE_TOOL containers mount *above* the call widget, so they
            # must not become the anchor — the result still needs to land
            # after the call widget, not after the hook container.
            if event.scope != HookType.BEFORE_TOOL and event.tool_call_id:
                self._tool_call_anchors[event.tool_call_id] = container
        else:
            await container.remove()

    async def handle_event(  # noqa: PLR0912
        self, event: BaseEvent, loading_widget: LoadingWidget | None = None
    ) -> ToolCallMessage | None:
        match event:
            case ReasoningEvent():
                await self._handle_reasoning_message(event)
            case AssistantEvent():
                await self._handle_assistant_message(event)
            case ToolCallEvent():
                await self.finalize_streaming()
                return await self._handle_tool_call(event, loading_widget)
            case ToolResultEvent():
                await self.finalize_streaming()
                sanitized_event = self._sanitize_event(event)
                await self._handle_tool_result(sanitized_event)
            case ToolStreamEvent():
                await self._handle_tool_stream(event)
            case CompactStartEvent():
                await self.finalize_streaming()
                await self._handle_compact_start()
            case CompactEndEvent():
                await self.finalize_streaming()
                await self._handle_compact_end(event)
            case AgentProfileChangedEvent():
                if self.on_profile_changed:
                    self.on_profile_changed()
            case SessionTitleUpdatedEvent():
                pass
            case UserMessageEvent():
                await self.finalize_streaming()
            case HookEvent():
                await self._handle_hook_event(event, loading_widget)
            case PlanReviewRequestedEvent():
                await self._handle_start_plan_review(file_path=event.file_path)
            case PlanReviewEndedEvent():
                self._handle_stop_plan_review()
            case WaitingForInputEvent():
                await self.finalize_streaming()
            case _:
                await self.finalize_streaming()
                await self._handle_unknown_event(event)
        return None

    def _sanitize_event(self, event: ToolResultEvent) -> ToolResultEvent:
        if isinstance(event, ToolResultEvent):
            return ToolResultEvent(
                tool_name=event.tool_name,
                tool_class=event.tool_class,
                result=event.result,
                error=TaggedText.from_string(event.error).message
                if event.error
                else None,
                skipped=event.skipped,
                skip_reason=TaggedText.from_string(event.skip_reason).message
                if event.skip_reason
                else None,
                cancelled=event.cancelled,
                duration=event.duration,
                tool_call_id=event.tool_call_id,
            )
        return event

    async def _handle_tool_call(
        self, event: ToolCallEvent, loading_widget: LoadingWidget | None = None
    ) -> ToolCallMessage | None:
        tool_call_id = event.tool_call_id
        existing_tool_call = self.tool_calls.get(tool_call_id) if tool_call_id else None
        if existing_tool_call:
            existing_tool_call.update_event(event)
            tool_call = existing_tool_call
        else:
            # A follow-up tool call is the recovery signal: leave prior errors muted.
            self._resolve_pending_errors(escalate=False)
            tool_call = ToolCallMessage(event)
            if tool_call_id:
                self.tool_calls[tool_call_id] = tool_call
                self._tool_call_anchors[tool_call_id] = tool_call
            await self.mount_callback(tool_call)

        if loading_widget and event.tool_class:
            adapter = ToolUIDataAdapter(event.tool_class)
            loading_widget.set_status(adapter.get_status_text())

        return tool_call

    async def _handle_tool_result(self, event: ToolResultEvent) -> None:
        tool_call_id = event.tool_call_id
        call_widget = self.tool_calls.get(tool_call_id) if tool_call_id else None
        anchor = (
            self._tool_call_anchors.get(tool_call_id) if tool_call_id else None
        ) or call_widget

        tool_result = ToolResultMessage(event, call_widget)
        await self.mount_callback(tool_result, after=anchor)

        if event.error and not event.skipped and not event.cancelled:
            self._pending_error_results.append(tool_result)

        if tool_call_id:
            self._tool_call_anchors[tool_call_id] = tool_result
            if tool_call_id in self.tool_calls:
                del self.tool_calls[tool_call_id]

    def _resolve_pending_errors(self, *, escalate: bool) -> None:
        if escalate:
            for result in self._pending_error_results:
                result.escalate_error()
        self._pending_error_results.clear()

    def escalate_unresolved_errors(self) -> None:
        # Called at turn end: errors not followed by a tool call were terminal.
        self._resolve_pending_errors(escalate=True)

    async def _handle_tool_stream(self, event: ToolStreamEvent) -> None:
        tool_call = self.tool_calls.get(event.tool_call_id)
        if tool_call:
            tool_call.set_stream_message(event.message)

    async def _handle_assistant_message(self, event: AssistantEvent) -> None:
        if self.current_streaming_reasoning is not None:
            self.current_streaming_reasoning.stop_spinning()
            await self.current_streaming_reasoning.stop_stream()
            self.current_streaming_reasoning = None

        if self.current_streaming_message is None:
            msg = AssistantMessage(event.content)
            self.current_streaming_message = msg
            await self.mount_callback(msg)
        else:
            await self.current_streaming_message.append_content(event.content)

    async def _handle_reasoning_message(self, event: ReasoningEvent) -> None:
        if self.current_streaming_message is not None:
            await self.current_streaming_message.stop_stream()
            if self.current_streaming_message.is_stripped_content_empty():
                await self.current_streaming_message.remove()
            self.current_streaming_message = None

        if self.current_streaming_reasoning is None:
            tools_collapsed = self.get_tools_collapsed()
            msg = ReasoningMessage(event.content, collapsed=tools_collapsed)
            self.current_streaming_reasoning = msg
            await self.mount_callback(msg)
        else:
            await self.current_streaming_reasoning.append_content(event.content)

    async def _handle_compact_start(self) -> None:
        compact_msg = CompactMessage()
        self.current_compact = compact_msg
        await self.mount_callback(compact_msg)

    async def _handle_compact_end(self, event: CompactEndEvent) -> None:
        if self.current_compact:
            self.current_compact.set_complete(
                old_session_id=event.old_session_id, new_session_id=event.new_session_id
            )
            self.current_compact = None

    async def _handle_unknown_event(self, event: BaseEvent) -> None:
        await self.mount_callback(NoMarkupStatic(str(event), classes="unknown-event"))

    async def finalize_streaming(self) -> None:
        if self.current_streaming_reasoning is not None:
            self.current_streaming_reasoning.stop_spinning()
            await self.current_streaming_reasoning.stop_stream()
            self.current_streaming_reasoning = None
        if self.current_streaming_message is not None:
            await self.current_streaming_message.stop_stream()
            self.current_streaming_message = None

    def stop_current_tool_call(
        self, success: bool = True, *, cancelled: bool = False
    ) -> None:
        for tool_call in self.tool_calls.values():
            if cancelled:
                # A user interrupt is neutral, not a failure: hold a grey square.
                tool_call.show_muted()
            else:
                tool_call.stop_spinning(success=success)
        self.tool_calls.clear()
        self._tool_call_anchors.clear()
        self._hook_containers.clear()
        # On cancel nothing is terminal -- leave prior errors muted too.
        self._resolve_pending_errors(escalate=not cancelled)

    def stop_current_compact(self) -> None:
        if self.current_compact:
            self.current_compact.stop_spinning(success=False)
            self.current_compact = None

    async def _handle_start_plan_review(self, file_path: Path) -> None:
        file_path.touch()
        msg = PlanFileMessage(file_path=file_path)
        self.plan_file_message = msg
        await self.mount_callback(msg)

    def _handle_stop_plan_review(self) -> None:
        if self.plan_file_message is None:
            return

        self.plan_file_message.stop_watching()
        self.plan_file_message = None