"""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