70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
"""In-process event bus: what the gateway does, as a stream frontends can tap.
|
|
|
|
Events are plain dicts ``{"type", "seq", "ts", ...}``. Publishers never
|
|
block; a subscriber that falls behind loses its oldest events rather than
|
|
stalling the turn that produced them. ``/api/events`` serialises the
|
|
stream as SSE, the panel builds its subagent tree from ``stream`` events
|
|
carrying ``parent_tool_use_id``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import itertools
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
__all__ = ["Event", "EventBus"]
|
|
|
|
Event = dict[str, Any]
|
|
|
|
|
|
class EventBus:
|
|
def __init__(self, *, maxsize: int = 2000) -> None:
|
|
self._maxsize = maxsize
|
|
self._subscribers: set[asyncio.Queue[Event]] = set()
|
|
self._seq = itertools.count(1)
|
|
|
|
def publish(self, type_: str, **data: Any) -> Event:
|
|
event: Event = {
|
|
"type": type_,
|
|
"seq": next(self._seq),
|
|
"ts": datetime.now(UTC).isoformat(timespec="milliseconds"),
|
|
**data,
|
|
}
|
|
for queue in list(self._subscribers):
|
|
if queue.full():
|
|
with contextlib.suppress(asyncio.QueueEmpty):
|
|
queue.get_nowait()
|
|
queue.put_nowait(event)
|
|
return event
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def subscribe(self) -> AsyncIterator[asyncio.Queue[Event]]:
|
|
queue: asyncio.Queue[Event] = asyncio.Queue(maxsize=self._maxsize)
|
|
self._subscribers.add(queue)
|
|
try:
|
|
yield queue
|
|
finally:
|
|
self._subscribers.discard(queue)
|
|
|
|
async def stream(
|
|
self, *, conversation_id: str | None = None
|
|
) -> AsyncIterator[Event]:
|
|
async with self.subscribe() as queue:
|
|
while True:
|
|
event = await queue.get()
|
|
if (
|
|
conversation_id is None
|
|
or event.get("conversation_id") == conversation_id
|
|
):
|
|
yield event
|
|
|
|
@property
|
|
def subscribers(self) -> int:
|
|
return len(self._subscribers)
|