beaver-gateway/src/beaver_gateway/conversations/state.py

"""What every layer of the conversations service shares: wiring and lookups."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any, cast

from claude_agent_sdk import project_key_for_directory

from beaver_gateway.conversations.injects import InjectQueue
from beaver_gateway.conversations.texts import ConversationTexts

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable, Sequence

    from claude_agent_sdk import SessionStore

    from beaver_gateway.agents.claude import ClaudeAgent
    from beaver_gateway.app import AgentRegistry
    from beaver_gateway.backends.claude_sdk import ClaudeSdkBackend
    from beaver_gateway.backends.sessions import SessionPool
    from beaver_gateway.conversations.distill import Distiller
    from beaver_gateway.conversations.envelope import Envelope
    from beaver_gateway.conversations.texts import UserSaid
    from beaver_gateway.events.bus import EventBus
    from beaver_gateway.frontends.base import Frontend
    from beaver_gateway.jobs.scheduler import Scheduler
    from beaver_gateway.storage.db import Database
    from beaver_gateway.storage.models import Conversation

__all__ = ["Question", "Runner", "State", "aware", "iso"]


@dataclass
class Runner:
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    wake: asyncio.Event = field(default_factory=asyncio.Event)
    task: asyncio.Task[None] | None = None
    turn_id: str | None = None
    origin: str | None = None
    text: str | None = None
    started_at: datetime | None = None
    tools: dict[str, dict[str, Any]] = field(default_factory=dict)

    def snapshot(self) -> dict[str, Any] | None:
        if self.turn_id is None:
            return None
        return {
            "id": self.turn_id,
            "origin": self.origin,
            "text": self.text,
            "started_at": iso(self.started_at),
            "tools": list(self.tools.values()),
        }


@dataclass(frozen=True, slots=True)
class Question:
    conversation_id: str
    turn_id: str | None
    questions: list[dict[str, Any]]
    answer: asyncio.Future[str]


class State:
    _db: Database
    _agents: AgentRegistry
    _backends: dict[str, Any]
    _bus: EventBus
    _pool: SessionPool
    _store: SessionStore
    _texts: ConversationTexts
    _frontends: list[Frontend]
    _normal_window: float
    _idle_days: tuple[int, ...]
    _idle_interval: float
    _question_timeout: float
    _envelope: Envelope | None
    _distiller: Distiller | None
    _user_sink: Callable[[UserSaid], Awaitable[None] | None] | None
    _queue: InjectQueue
    _runners: dict[int, Runner]
    _questions: dict[str, Question]
    _tasks: set[asyncio.Task[None]]
    _idle_task: asyncio.Task[None] | None
    scheduler: Scheduler | None

    def __init__(
        self,
        *,
        db: Database,
        agents: AgentRegistry,
        backends: dict[str, Any],
        bus: EventBus,
        pool: SessionPool,
        store: SessionStore,
        texts: ConversationTexts | None = None,
        frontends: Sequence[Frontend] = (),
        normal_window: float = 3600.0,
        idle_days: Sequence[int] = (2,),
        idle_interval: float = 3600.0,
        question_timeout: float = 600.0,
        envelope: Envelope | None = None,
        distiller: Distiller | None = None,
        user_sink: Callable[[UserSaid], Awaitable[None] | None] | None = None,
    ) -> None:
        self._db = db
        self._agents = agents
        self._backends = backends
        self._bus = bus
        self._pool = pool
        self._store = store
        self._texts = texts or ConversationTexts()
        self._frontends = [f for f in frontends if f.name]
        self._normal_window = normal_window
        self._idle_days = tuple(sorted(idle_days))
        self._idle_interval = idle_interval
        self._question_timeout = question_timeout
        self._envelope = envelope
        self._distiller = distiller
        self._user_sink = user_sink
        self._queue = InjectQueue(db)
        self._runners = {}
        self._questions = {}
        self._tasks = set()
        self._idle_task = None
        self.scheduler = None
        for backend in backends.values():
            if hasattr(backend, "waker"):
                backend.waker = self._woken

    async def _woken(self, key: str, note: str) -> None:
        raise NotImplementedError

    @property
    def db(self) -> Database:
        return self._db

    @property
    def queue(self) -> InjectQueue:
        return self._queue

    @property
    def bus(self) -> EventBus:
        return self._bus

    @property
    def pool(self) -> SessionPool:
        return self._pool

    def _backend(self, agent: str) -> ClaudeSdkBackend:
        backend = self._backends.get(agent)
        if backend is None or not hasattr(backend, "repair_session"):
            msg = f"agent {agent!r} has no Claude SDK backend"
            raise LookupError(msg)
        return cast("ClaudeSdkBackend", backend)

    def _claude_agent(self, name: str) -> ClaudeAgent:
        agent = self._agents.get(name)
        if agent is None or not hasattr(agent, "cwd"):
            msg = f"unknown Claude agent {name!r}"
            raise LookupError(msg)
        return cast("ClaudeAgent", agent)

    def _store_key(self, conv: Conversation) -> dict[str, str]:
        agent = self._claude_agent(conv.agent_name)
        return {
            "project_key": project_key_for_directory(str(agent.cwd)),
            "session_id": cast("str", conv.session_id),
        }

    def _runner(self, row_id: int) -> Runner:
        runner = self._runners.get(row_id)
        if runner is None:
            runner = Runner()
            self._runners[row_id] = runner
        return runner

    def _track(self, task: asyncio.Task[None]) -> None:
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)


def aware(value: datetime) -> datetime:
    return value if value.tzinfo is not None else value.replace(tzinfo=UTC)


def iso(value: datetime | None) -> str | None:
    return aware(value).isoformat(timespec="seconds") if value is not None else None