smart-org/apps/api/app/application/engine/events.py

"""The engine's live feed: batched ``run.log`` / ``run.tool`` / ``task.gate``.

PLAN-v2 §5 requires log events to be *batched* — at most one event per 500 ms
per stream, at most 200 lines per event. A model loop that emits one WebSocket
event per tool call turns a chatty agent into a flood that buries the browser,
so the rule is not a nicety.

``RunFeed`` buffers log lines per run and publishes them at the boundaries the
runner chooses: between steps (rate-limited) and at the end of a turn (always).
Tool results and gate verdicts are their own event types and never wait behind
the log buffer — a command's output is the thing the live log panel exists for.

Publishing goes through the generic envelope rather than a typed emitter
because ``ports/realtime.py`` is a frozen contract and these three event types
have no typed helper there.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field

from app.application.ports.realtime import (
    RealtimeEvent,
    org_channel,
    publish,
    run_channel,
    task_channel,
)
from app.core.config import settings


@dataclass(slots=True, kw_only=True)
class _Stream:
    """One run's buffer: the lines waiting, and when we last published."""

    run_id: str
    task_id: str
    organization_id: str
    lines: list[str] = field(default_factory=list)
    last_publish: float = 0.0
    #: Set when the buffer reached its line cap, so the next flush must not
    #: wait for the interval to elapse.
    overflow: bool = False


class RunFeed:
    """Batches engine output into bounded realtime events, per run."""

    def __init__(
        self, *, min_interval: float | None = None, max_lines: int | None = None
    ) -> None:
        self._min_interval = (
            settings.engine_log_flush_seconds if min_interval is None else min_interval
        )
        self._max_lines = (
            settings.engine_log_max_lines_per_event if max_lines is None else max_lines
        )
        self._streams: dict[str, _Stream] = {}
        #: Events published, by type — so a test can assert "one batched log
        #: event, not one per line" without touching the realtime bus.
        self.published: dict[str, int] = {}

    # -- buffering ----------------------------------------------------------

    def log(self, *, run_id: str, task_id: str, organization_id: str, line: str) -> None:
        """Buffer one log line. Nothing is published until a flush."""
        stream = self._stream(run_id, task_id, organization_id)
        stream.lines.append(line)
        if len(stream.lines) >= self._max_lines:
            # The buffer is full: the next flush must not wait for the timer.
            stream.overflow = True

    async def tool(
        self,
        *,
        run_id: str,
        task_id: str,
        organization_id: str,
        tool: str,
        ok: bool,
        summary: str,
    ) -> None:
        """Report one executed tool call as a single, self-contained event."""
        await self._emit(
            RealtimeEvent(
                type="run.tool",
                channels=(
                    run_channel(run_id),
                    task_channel(task_id),
                    org_channel(organization_id),
                ),
                organization_id=organization_id,
                data={
                    "run_id": run_id,
                    "task_id": task_id,
                    "tool": tool,
                    "ok": ok,
                    "summary": summary[:2000],
                },
            )
        )

    async def gate(
        self,
        *,
        task_id: str,
        organization_id: str,
        gate: str,
        ok: bool,
        detail: str,
    ) -> None:
        """Report one gate verdict, with the evidence attached."""
        await self._emit(
            RealtimeEvent(
                type="task.gate",
                channels=(task_channel(task_id), org_channel(organization_id)),
                organization_id=organization_id,
                data={
                    "task_id": task_id,
                    "gate": gate,
                    "ok": ok,
                    "detail": detail[:8000],
                },
            )
        )

    # -- flushing -----------------------------------------------------------

    async def flush(self, *, run_id: str, force: bool = False) -> None:
        """Publish a run's buffered lines, unless the interval has not elapsed.

        ``force`` is for the end of a turn: whatever is buffered then is the
        last thing the browser will see until the next turn, so it goes out.
        """
        stream = self._streams.get(run_id)
        if stream is None or not stream.lines:
            return
        now = time.monotonic()
        elapsed = now - stream.last_publish
        if not force and not stream.overflow and elapsed < self._min_interval:
            return
        await self._publish(stream, force=force or stream.overflow)

    async def flush_all(self) -> None:
        for stream in list(self._streams.values()):
            await self._publish(stream, force=True)

    # -- internals ----------------------------------------------------------

    async def _publish(self, stream: _Stream, *, force: bool) -> None:
        if force:
            lines, stream.lines = list(stream.lines), []
        else:
            lines, stream.lines = stream.lines[: self._max_lines], stream.lines[self._max_lines :]
        stream.overflow = False
        if not lines:
            return
        stream.last_publish = time.monotonic()
        await self._emit(
            RealtimeEvent(
                type="run.log",
                channels=(
                    run_channel(stream.run_id),
                    task_channel(stream.task_id),
                    org_channel(stream.organization_id),
                ),
                organization_id=stream.organization_id,
                data={"run_id": stream.run_id, "task_id": stream.task_id, "lines": lines},
            )
        )

    async def _emit(self, event: RealtimeEvent) -> None:
        self.published[event.type] = self.published.get(event.type, 0) + 1
        # The realtime publisher is documented as never raising; a counter
        # keeps tests honest even when the bus is a NullPublisher.
        await publish(event)

    def _stream(self, run_id: str, task_id: str, organization_id: str) -> _Stream:
        stream = self._streams.get(run_id)
        if stream is None:
            stream = _Stream(
                run_id=run_id, task_id=task_id, organization_id=organization_id
            )
            self._streams[run_id] = stream
        return stream


__all__ = ["RunFeed"]