refactor: split flat core into capability packages, layer the conversations service, English defaults for every model-facing text

src/beaver_gateway/conversations/turns.py

@@ -0,0 +1,444 @@
+"""Running turns: one worker per conversation over its queue, the backend call."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import inspect
+import logging
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any, cast
+from uuid import uuid4
+
+from claude_agent_sdk import (
+    AssistantMessage,
+    RateLimitEvent,
+    ResultMessage,
+    StreamEvent,
+    ToolResultBlock,
+    ToolUseBlock,
+    UserMessage,
+)
+
+from beaver_gateway.backends.capture import TurnCapture
+from beaver_gateway.backends.transcript import text_of
+from beaver_gateway.conversations import injects
+from beaver_gateway.conversations.seeds import Seeds
+from beaver_gateway.conversations.state import Runner, aware, iso
+from beaver_gateway.conversations.texts import UserSaid
+from beaver_gateway.frontends.accumulate import StreamAccumulator
+from beaver_gateway.storage.models import Conversation, InjectQueueItem, RateLimit
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable, Sequence
+
+    from beaver_gateway.events.stream import MessageStreamEvent
+
+__all__ = ["Turns"]
+
+_log = logging.getLogger(__name__)
+
+
+class Turns(Seeds):
+    async def turn(
+        self,
+        conv: Conversation,
+        *,
+        messages: Sequence[Any],
+        origin: str,
+        capture: TurnCapture | None = None,
+        session_id: str | None = None,
+        use_session: bool = True,
+        tools: bool = True,
+        turn_id: str | None = None,
+        item_origin: str | None = None,
+    ) -> AsyncIterator[MessageStreamEvent]:
+        """Run one turn under the conversation's lock; the only path to the backend."""
+        row_id = cast("int", conv.id)
+        runner = self._runner(row_id)
+        backend = self._backend(conv.agent_name)
+        turn_id = turn_id or f"turn_{uuid4().hex[:12]}"
+        capture = capture or TurnCapture()
+        resume = session_id if session_id is not None else conv.session_id
+        async with runner.lock:
+            runner.turn_id = turn_id
+            runner.origin = origin
+            runner.text = _prompt_preview(messages)
+            runner.started_at = datetime.now(UTC)
+            runner.tools = {}
+            await self._mark_running(conv, turn_id)
+            before = await self.before_turn(conv)
+            self._bus.publish(
+                "turn.start",
+                conversation_id=conv.external_id,
+                turn_id=turn_id,
+                origin=origin,
+                item_origin=item_origin,
+                text=runner.text,
+            )
+            stop = "error"
+            cut = False
+            try:
+                events = backend.complete(
+                    agent=self._claude_agent(conv.agent_name),
+                    messages=messages,
+                    conversation_id=conv.external_id,
+                    session_id=resume if use_session else None,
+                    reseed=not use_session,
+                    capture=capture,
+                    kind=conv.kind,
+                    pinned=conv.kind == "master",
+                    tools=tools,
+                    observer=self._observer(conv, runner, turn_id, origin),
+                    turn_id=turn_id,
+                )
+                async for event in events:
+                    yield event
+                stop = "interrupted" if capture.interrupted else "end_turn"
+            except asyncio.CancelledError:
+                cut = True
+                raise
+            finally:
+                runner.turn_id = None
+                await self._mark_done(conv, capture, cut=cut)
+                if stop != "error":
+                    try:
+                        await self.after_turn(conv, before)
+                    except Exception:  # noqa: BLE001
+                        _log.exception("after-turn hook on %s failed", conv.external_id)
+                self._bus.publish(
+                    "turn.end",
+                    conversation_id=conv.external_id,
+                    turn_id=turn_id,
+                    origin=origin,
+                    item_origin=item_origin,
+                    stop=stop,
+                    usage=_usage_dict(capture),
+                )
+
+    async def run_text_turn(
+        self,
+        conv: Conversation,
+        text: str,
+        *,
+        origin: str,
+        tools: bool = True,
+        turn_id: str | None = None,
+        item_origin: str | None = None,
+    ) -> tuple[str, TurnCapture]:
+        capture = TurnCapture()
+        acc = StreamAccumulator()
+        agent = self._claude_agent(conv.agent_name)
+        async for event in self.turn(
+            conv,
+            messages=[{"role": "user", "content": text}],
+            origin=origin,
+            capture=capture,
+            tools=tools,
+            turn_id=turn_id,
+            item_origin=item_origin,
+        ):
+            acc.feed(event)
+        message = acc.finalize(model=agent.model)
+        reply = "\n\n".join(
+            getattr(b, "text", "")
+            for b in message.content
+            if getattr(b, "type", "") == "text"
+        ).strip()
+        return reply, capture
+
+    async def before_turn(self, conv: Conversation) -> str | None:  # noqa: ARG002
+        return None
+
+    async def after_turn(self, conv: Conversation, before: str | None) -> None:  # noqa: ARG002
+        return
+
+    def turn_origin(self, conv: Conversation) -> str | None:
+        runner = self._runners.get(cast("int", conv.id))
+        return runner.origin if runner is not None and runner.turn_id else None
+
+    def _ensure_worker(self, row_id: int) -> None:
+        runner = self._runner(row_id)
+        runner.wake.set()
+        if runner.task is None or runner.task.done():
+            runner.task = asyncio.create_task(self._worker(row_id))
+            self._track(runner.task)
+
+    async def _worker(self, row_id: int) -> None:
+        runner = self._runner(row_id)
+        while True:
+            items = await self._queue.pending(row_id)
+            batch, wait = self._pick(items)
+            if batch is None:
+                runner.wake.clear()
+                if wait is None:
+                    return
+                with contextlib.suppress(TimeoutError):
+                    await asyncio.wait_for(runner.wake.wait(), timeout=wait)
+                continue
+            conv = await self.get_row(row_id)
+            if conv is None:
+                await self._queue.finish(batch, status="failed")
+                return
+            await self._run_batch(conv, batch)
+
+    def _pick(
+        self, items: list[InjectQueueItem]
+    ) -> tuple[list[InjectQueueItem] | None, float | None]:
+        if not items:
+            return None, None
+        head = items[0]
+        if head.priority == "urgent":
+            return [head], None
+        tail = [i for i in items if i is not head and i.priority in ("wake", "normal")]
+        if head.priority in ("user", "wake"):
+            return [head, *tail], None
+        age = (datetime.now(UTC) - aware(head.created_at)).total_seconds()
+        if age >= self._normal_window:
+            return [head, *tail], None
+        return None, max(self._normal_window - age, 1.0)
+
+    async def _run_batch(
+        self, conv: Conversation, batch: list[InjectQueueItem]
+    ) -> None:
+        head = batch[0]
+        turn_id = f"turn_{uuid4().hex[:12]}"
+        await self._queue.start(batch, turn_id)
+        if head.priority == "user":
+            origin = "user"
+            prompt = head.text
+            await self._note_user(conv, head.text)
+            envelope = self._envelope_for(conv, head.text)
+            if envelope:
+                prompt += "\n\n" + envelope
+            if len(batch) > 1:
+                prompt += "\n\n" + self._bundle(batch[1:])
+        else:
+            origin = "inject"
+            prompt = "\n\n".join(
+                f"{self._texts.inject_header(injects.context_of(i))}\n{i.text}"
+                for i in batch
+            )
+        seed = await self.pending_seed(conv)
+        if seed:
+            prompt = f"{seed}\n\n{prompt}"
+        try:
+            text, capture = await self.run_text_turn(
+                conv, prompt, origin=origin, turn_id=turn_id, item_origin=head.origin
+            )
+        except Exception:  # noqa: BLE001
+            _log.exception("turn %s on %s failed", turn_id, conv.external_id)
+            await self._queue.finish(batch, status="failed")
+            return
+        await self._queue.finish(
+            batch, status="interrupted" if capture.interrupted else "done"
+        )
+        if origin == "user":
+            self._bus.publish(
+                "reply",
+                conversation_id=conv.external_id,
+                turn_id=turn_id,
+                item=head.id,
+                item_origin=head.origin,
+                source="queue",
+                prompt=prompt,
+                user_text=head.text,
+                text=text,
+            )
+
+    def _bundle(self, items: Sequence[InjectQueueItem]) -> str:
+        lines = [self._texts.bundle_header.format(since=iso(items[0].created_at))]
+        lines.extend(f"- [{i.origin}] {i.text}" for i in items)
+        return "\n".join(lines)
+
+    def _envelope_for(self, conv: Conversation, text: str = "") -> str | None:
+        if self._envelope is None:
+            return None
+        if conv.kind == "master":
+            return self._envelope.build(text=text, kind="master")
+        if conv.kind == "branch":
+            return self._envelope.recall_only(text=text, kind="branch")
+        return None
+
+    async def _note_user(self, conv: Conversation, text: str) -> None:
+        if self._user_sink is None:
+            return
+        message = UserSaid(
+            conversation_id=conv.external_id,
+            kind=conv.kind,
+            title=conv.title,
+            text=text,
+            at=datetime.now(UTC),
+        )
+        try:
+            result = self._user_sink(message)
+            if inspect.isawaitable(result):
+                await result
+        except Exception:  # noqa: BLE001
+            _log.exception("user sink failed for %s", conv.external_id)
+
+    def _observer(
+        self, conv: Conversation, runner: Runner, turn_id: str, origin: str
+    ) -> Callable[[Any], None]:
+        conversation_id = conv.external_id
+
+        def observe(message: Any) -> None:
+            parent = getattr(message, "parent_tool_use_id", None)
+            if isinstance(message, RateLimitEvent):
+                self._observe_rate_limit(conv, message)
+            elif isinstance(message, StreamEvent):
+                self._bus.publish(
+                    "stream",
+                    conversation_id=conversation_id,
+                    turn_id=turn_id,
+                    origin=origin,
+                    parent_tool_use_id=parent,
+                    event=message.event,
+                )
+            elif isinstance(message, AssistantMessage):
+                for block in message.content:
+                    if isinstance(block, ToolUseBlock):
+                        event = self._bus.publish(
+                            "tool",
+                            conversation_id=conversation_id,
+                            turn_id=turn_id,
+                            origin=origin,
+                            parent_tool_use_id=parent,
+                            tool_use_id=block.id,
+                            name=block.name,
+                            input=block.input,
+                        )
+                        runner.tools[block.id] = {
+                            "tool_use_id": block.id,
+                            "name": block.name,
+                            "input": block.input,
+                            "parent_tool_use_id": parent,
+                            "started_at": event["ts"],
+                            "ended_at": None,
+                            "is_error": None,
+                            "content": None,
+                        }
+            elif isinstance(message, UserMessage):
+                blocks = message.content if isinstance(message.content, list) else ()
+                for block in blocks:
+                    if isinstance(block, ToolResultBlock):
+                        event = self._bus.publish(
+                            "tool.result",
+                            conversation_id=conversation_id,
+                            turn_id=turn_id,
+                            origin=origin,
+                            parent_tool_use_id=parent,
+                            tool_use_id=block.tool_use_id,
+                            is_error=bool(block.is_error),
+                            content=_result_preview(block.content),
+                        )
+                        node = runner.tools.get(block.tool_use_id)
+                        if node is not None:
+                            node["ended_at"] = event["ts"]
+                            node["is_error"] = event["is_error"]
+                            node["content"] = event["content"]
+            elif isinstance(message, ResultMessage) and parent is None:
+                self._bus.publish(
+                    "result",
+                    conversation_id=conversation_id,
+                    turn_id=turn_id,
+                    origin=origin,
+                    subtype=message.subtype,
+                    is_error=message.is_error,
+                    num_turns=message.num_turns,
+                )
+
+        return observe
+
+    def _observe_rate_limit(self, conv: Conversation, message: RateLimitEvent) -> None:
+        info = message.rate_limit_info
+        row = RateLimit(
+            window=info.rate_limit_type or "unknown",
+            status=info.status,
+            utilization=info.utilization,
+            resets_at=_from_unix(info.resets_at),
+            overage_status=info.overage_status,
+            overage_resets_at=_from_unix(info.overage_resets_at),
+            agent_name=conv.agent_name,
+            session_id=message.session_id,
+            raw=dict(info.raw),
+        )
+        self._bus.publish(
+            "rate_limit",
+            conversation_id=conv.external_id,
+            window=row.window,
+            status=row.status,
+            utilization=row.utilization,
+            resets_at=iso(row.resets_at),
+            overage_status=row.overage_status,
+        )
+        self._track(asyncio.create_task(self._record_rate_limit(row)))
+
+    async def _record_rate_limit(self, row: RateLimit) -> None:
+        try:
+            async with self._db.session() as session:
+                session.add(row)
+                await session.commit()
+        except Exception:  # noqa: BLE001
+            _log.exception("rate limit write failed")
+
+    async def _mark_running(self, conv: Conversation, turn_id: str) -> None:
+        async def apply(row: Conversation) -> None:
+            row.running_turn = turn_id
+            row.last_activity_at = datetime.now(UTC)
+
+        await self._update(conv, apply)
+
+    async def _mark_done(
+        self, conv: Conversation, capture: TurnCapture, *, cut: bool = False
+    ) -> None:
+        async def apply(row: Conversation) -> None:
+            if not cut:
+                row.running_turn = None
+            row.last_activity_at = datetime.now(UTC)
+            if capture.session_id is not None:
+                row.session_id = capture.session_id
+
+        await self._update(conv, apply)
+
+
+def _prompt_preview(messages: Sequence[Any], limit: int = 400) -> str | None:
+    if not messages:
+        return None
+    text = text_of(messages[-1].get("content"))
+    return text[:limit] if text else None
+
+
+def _from_unix(value: int | None) -> datetime | None:
+    return datetime.fromtimestamp(value, tz=UTC) if value is not None else None
+
+
+def _result_preview(
+    content: str | list[dict[str, Any]] | None, limit: int = 400
+) -> str:
+    if content is None:
+        return ""
+    text = (
+        content
+        if isinstance(content, str)
+        else "\n".join(
+            str(part.get("text", ""))
+            for part in content
+            if isinstance(part, dict) and part.get("type") == "text"
+        )
+    )
+    return text if len(text) <= limit else text[:limit] + "…"
+
+
+def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None:
+    usage = capture.usage
+    if usage is None:
+        return None
+    return {
+        "input": usage.input_tokens,
+        "output": usage.output_tokens,
+        "cache_read": usage.cache_read_tokens,
+        "cache_creation": usage.cache_creation_tokens,
+        "cost_usd": usage.cost_usd,
+        "duration_ms": usage.duration_ms,
+    }