hkuds--deeptutor/deeptutor/core/stream_bus.py

"""
Stream Bus
==========

Async event channel that tools / capabilities emit into and consumers
(CLI renderer, WebSocket pusher, JSON writer) read from.

Usage::

    bus = StreamBus()

    # Producer side (inside a capability)
    await bus.emit(StreamEvent(type=StreamEventType.CONTENT, content="Hello"))

    # Consumer side
    async for event in bus.subscribe():
        print(event.content)
"""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
import json
from typing import Any, AsyncIterator

from .stream import StreamEvent, StreamEventType
from .trace import merge_trace_metadata


class StreamBus:
    """Fan-out async event bus for a single chat turn."""

    def __init__(self) -> None:
        self._subscribers: list[asyncio.Queue[StreamEvent | None]] = []
        self._closed = False
        self._history: list[StreamEvent] = []
        self._input_listeners: list[asyncio.Queue[str]] = []

    async def emit(self, event: StreamEvent) -> None:
        """Push *event* to every active subscriber."""
        if self._closed:
            return
        self._history.append(event)
        for q in self._subscribers:
            await q.put(event)

    async def subscribe(self) -> AsyncIterator[StreamEvent]:
        """Yield events until the bus is closed."""
        q: asyncio.Queue[StreamEvent | None] = asyncio.Queue()
        self._subscribers.append(q)
        # Snapshot the replay range in the same synchronous step as the
        # queue registration: events emitted while we replay are delivered
        # via the queue only. Iterating the live list instead would yield
        # those events twice (list-append during iteration + queue copy).
        replay_count = len(self._history)
        try:
            for event in self._history[:replay_count]:
                yield event
            if self._closed and q.empty():
                return
            while True:
                event = await q.get()
                if event is None:
                    break
                yield event
        finally:
            self._subscribers.remove(q)

    async def close(self) -> None:
        """Signal all subscribers that the stream is finished."""
        self._closed = True
        for q in self._subscribers:
            await q.put(None)

    # ---- convenience helpers for producers ----

    @asynccontextmanager
    async def stage(
        self,
        name: str,
        source: str = "",
        metadata: dict[str, Any] | None = None,
    ):
        """Context manager that emits STAGE_START / STAGE_END around a block."""
        await self.emit(
            StreamEvent(
                type=StreamEventType.STAGE_START,
                source=source,
                stage=name,
                metadata=metadata or {},
            )
        )
        try:
            yield
        finally:
            await self.emit(
                StreamEvent(
                    type=StreamEventType.STAGE_END,
                    source=source,
                    stage=name,
                    metadata=metadata or {},
                )
            )

    async def content(
        self,
        text: str,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.CONTENT,
                source=source,
                stage=stage,
                content=text,
                metadata=metadata or {},
            )
        )

    async def thinking(
        self,
        text: str,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.THINKING,
                source=source,
                stage=stage,
                content=text,
                metadata=metadata or {},
            )
        )

    async def observation(
        self,
        text: str,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.OBSERVATION,
                source=source,
                stage=stage,
                content=text,
                metadata=metadata or {},
            )
        )

    async def tool_call(
        self,
        tool_name: str,
        args: dict[str, Any],
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.TOOL_CALL,
                source=source,
                stage=stage,
                content=tool_name,
                metadata=merge_trace_metadata({"args": args}, metadata),
            )
        )

    async def tool_result(
        self,
        tool_name: str,
        result: str,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.TOOL_RESULT,
                source=source,
                stage=stage,
                content=result,
                metadata=merge_trace_metadata({"tool": tool_name}, metadata),
            )
        )

    async def progress(
        self,
        message: str,
        current: int = 0,
        total: int = 0,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.PROGRESS,
                source=source,
                stage=stage,
                content=message,
                metadata=merge_trace_metadata(
                    {"current": current, "total": total},
                    metadata,
                ),
            )
        )

    async def sources(
        self,
        sources: list[dict[str, Any]],
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.SOURCES,
                source=source,
                stage=stage,
                metadata=merge_trace_metadata({"sources": sources}, metadata),
            )
        )

    async def result(
        self,
        data: dict[str, Any],
        source: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.RESULT,
                source=source,
                metadata=merge_trace_metadata(data, metadata),
            )
        )

    async def error(
        self,
        message: str,
        source: str = "",
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.emit(
            StreamEvent(
                type=StreamEventType.ERROR,
                source=source,
                stage=stage,
                content=message,
                metadata=metadata or {},
            )
        )

    async def wait_for_input(
        self,
        prompt: str,
        source: str = "",
        stage: str = "",
        timeout: float | None = None,
    ) -> str:
        """Pause capability execution and wait for user input from the frontend.

        Returns the user's input, or an empty string if *timeout* seconds elapse
        with no input (e.g. when running from a CLI that cannot send input).
        Pass ``timeout=None`` (the default) to wait indefinitely for interactive
        clients; pass a finite value for headless/CLI entry points.
        """
        await self.emit(
            StreamEvent(
                type=StreamEventType.WAIT_FOR_INPUT,
                source=source,
                stage=stage,
                content=prompt,
            )
        )
        input_queue: asyncio.Queue[str] = asyncio.Queue()
        self._input_listeners.append(input_queue)
        try:
            return await asyncio.wait_for(input_queue.get(), timeout=timeout)
        except asyncio.TimeoutError:
            return ""
        finally:
            if input_queue in self._input_listeners:
                self._input_listeners.remove(input_queue)

    def submit_input(self, content: str) -> None:
        """Receive user input from the frontend/WS and deliver to waiters."""
        for q in self._input_listeners:
            q.put_nowait(content)
        self._input_listeners.clear()

    # ---- consumer adapters ----

    @staticmethod
    def event_to_json(event: StreamEvent) -> str:
        """Serialize an event to a single-line JSON string (NDJSON)."""
        return json.dumps(event.to_dict(), ensure_ascii=False)


# ---- per-turn bus registry for user_input routing ----

_bus_registry: dict[str, StreamBus] = {}


def register_bus(turn_id: str, bus: StreamBus) -> None:
    """Register a bus so ``user_input`` WS messages can find it."""
    _bus_registry[turn_id] = bus


def unregister_bus(turn_id: str) -> None:
    """Remove a bus from the registry."""
    _bus_registry.pop(turn_id, None)


def get_bus(turn_id: str) -> StreamBus | None:
    """Look up the active bus for *turn_id*."""
    return _bus_registry.get(turn_id)