atc-agents/api/agent_terminal.py

"""Per-agent live terminal buffers and streaming."""

from __future__ import annotations

import asyncio
import uuid
from collections import deque
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable

MAX_LINES_PER_AGENT = 300

PublishFn = Callable[[dict[str, Any]], Awaitable[None]]

_buffers: dict[str, deque[dict[str, Any]]] = {}
_publish: PublishFn | None = None
_loop: asyncio.AbstractEventLoop | None = None


def init_terminals(agent_ids: list[str]) -> None:
    for aid in agent_ids:
        if aid not in _buffers:
            _buffers[aid] = deque(maxlen=MAX_LINES_PER_AGENT)


def set_terminal_publisher(fn: PublishFn) -> None:
    global _publish, _loop
    _publish = fn
    try:  # capture the main event loop so background threads can stream too
        _loop = asyncio.get_running_loop()
    except RuntimeError:
        _loop = None


def _build_line(agent_id: str, text: str, level: str, phase: str, prompt_id: str | None) -> dict[str, Any]:
    return {
        "id": str(uuid.uuid4())[:8],
        "ts": datetime.now(timezone.utc).isoformat(),
        "agent_id": agent_id,
        "level": level,
        "phase": phase,
        "text": text,
        "prompt_id": prompt_id,
    }


def emit_threadsafe(agent_id: str, text: str, *, level: str = "info", phase: str = "ops") -> dict[str, Any]:
    """Append a terminal line and broadcast it from a non-async context (e.g. a
    background ``threading.Thread``). Safe to call from any thread."""
    init_terminals([agent_id])
    line = _build_line(agent_id, text, level, phase, None)
    _buffers[agent_id].append(line)
    if _publish and _loop and not _loop.is_closed():
        try:
            asyncio.run_coroutine_threadsafe(_publish({"type": "terminal", "line": line}), _loop)
        except Exception:
            pass
    return line


def get_terminal_lines(agent_id: str, limit: int = 200) -> list[dict[str, Any]]:
    buf = _buffers.get(agent_id, deque())
    items = list(buf)
    return items[-limit:]


def get_all_terminals(limit: int = 200) -> dict[str, list[dict[str, Any]]]:
    return {aid: get_terminal_lines(aid, limit) for aid in _buffers}


async def terminal_log(
    agent_id: str,
    text: str,
    *,
    level: str = "info",
    phase: str = "ops",
    prompt_id: str | None = None,
    mirror: bool = True,
) -> dict[str, Any]:
    init_terminals([agent_id])
    line = {
        "id": str(uuid.uuid4())[:8],
        "ts": datetime.now(timezone.utc).isoformat(),
        "agent_id": agent_id,
        "level": level,
        "phase": phase,
        "text": text,
        "prompt_id": prompt_id,
    }
    _buffers[agent_id].append(line)
    if _publish and mirror:
        await _publish({"type": "terminal", "line": line})
    return line


# Type: async (level, phase, text) -> None
TerminalLogFn = Callable[[str, str, str], Awaitable[None]]


def make_logger(agent_id: str, prompt_id: str | None = None) -> TerminalLogFn:
    async def log(level: str, phase: str, text: str) -> None:
        await terminal_log(agent_id, text, level=level, phase=phase, prompt_id=prompt_id)

    return log