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

src/beaver_gateway/app.py

@@ -0,0 +1,440 @@
+"""Build the runtime from a ``Gateway`` and run every part of it until shutdown."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+from contextlib import AsyncExitStack
+from typing import TYPE_CHECKING, Any
+
+import psycopg
+import uvicorn
+from pgqueuer import PsycopgDriver
+from raycast_api import Client as RaycastClient
+from raycast_api.config import Config as RaycastConfig
+
+from beaver_gateway.agents.claude import ClaudeAgent
+from beaver_gateway.agents.raycast import RaycastAgent
+from beaver_gateway.backends.claude_sdk import (
+    ClaudeSdkBackend,
+    RunnerConfig,
+    UsageEvent,
+)
+from beaver_gateway.backends.raycast import RaycastBackend
+from beaver_gateway.backends.sessions import SessionPool
+from beaver_gateway.conversations.envelope import Envelope
+from beaver_gateway.conversations.rotation import Rotation, RotationPolicy
+from beaver_gateway.conversations.service import Conversations
+from beaver_gateway.conversations.tools import build_tool_server
+from beaver_gateway.events.bus import EventBus
+from beaver_gateway.frontends.base import GatewayRuntime
+from beaver_gateway.frontends.bearer import require_token
+from beaver_gateway.frontends.root import build_root_app
+from beaver_gateway.jobs.scheduler import Scheduler
+from beaver_gateway.mcp.internal_app import build_internal_app
+from beaver_gateway.security.auth import TokenStore
+from beaver_gateway.storage import (
+    Database,
+    PostgresSessionStore,
+    Usage,
+    append_audit,
+    append_usage,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+
+    from claude_agent_sdk import McpSdkServerConfig
+    from fastmcp import FastMCP
+    from fastmcp.tools.base import Tool as FastMCPTool
+    from starlette.applications import Starlette
+    from starlette.requests import Request
+    from starlette.types import ASGIApp
+
+    from beaver_gateway.agents.base import BaseAgent
+    from beaver_gateway.agents.policy import ToolAudit
+    from beaver_gateway.backends.base import Backend
+    from beaver_gateway.config import Gateway
+    from beaver_gateway.mcp.types import McpServerT
+    from beaver_gateway.settings import Settings
+
+__all__ = ["AgentRegistry", "McpRegistry", "run"]
+
+_log = logging.getLogger("beaver_gateway.app")
+
+
+class AgentRegistry:
+    def __init__(self, agents: Iterable[BaseAgent]) -> None:
+        self._agents: dict[str, BaseAgent] = {}
+        for a in agents:
+            if a.name in self._agents:
+                msg = f"duplicate agent name: {a.name!r}"
+                raise ValueError(msg)
+            self._agents[a.name] = a
+
+    def __getitem__(self, name: str) -> BaseAgent:
+        return self._agents[name]
+
+    def get(self, name: str) -> BaseAgent | None:
+        return self._agents.get(name)
+
+    def __iter__(self) -> Iterator[BaseAgent]:
+        return iter(self._agents.values())
+
+    def __len__(self) -> int:
+        return len(self._agents)
+
+    def __contains__(self, name: object) -> bool:
+        return name in self._agents
+
+
+class McpRegistry:
+    def __init__(self, mcps: Iterable[McpServerT]) -> None:
+        self._mcps: dict[str, McpServerT] = {}
+        for m in mcps:
+            if m.name in self._mcps:
+                msg = f"duplicate mcp name: {m.name!r}"
+                raise ValueError(msg)
+            self._mcps[m.name] = m
+
+    def __getitem__(self, name: str) -> McpServerT:
+        return self._mcps[name]
+
+    def get(self, name: str) -> McpServerT | None:
+        return self._mcps.get(name)
+
+    def __iter__(self) -> Iterator[McpServerT]:
+        return iter(self._mcps.values())
+
+    def __len__(self) -> int:
+        return len(self._mcps)
+
+    def __contains__(self, name: object) -> bool:
+        return name in self._mcps
+
+
+async def run(gateway: Gateway, settings: Settings) -> None:
+    agents = AgentRegistry(gateway.agents)
+    mcps = McpRegistry(gateway.mcps)
+    db = Database(settings.database_url)
+    await db.create_all()
+    token_store = TokenStore(
+        db,
+        bootstrap=TokenStore.parse_bootstrap(settings.bootstrap_tokens),
+        bootstrap_scopes=TokenStore.parse_bootstrap_scopes(settings.bootstrap_tokens),
+    )
+
+    async with AsyncExitStack() as stack:
+        stack.push_async_callback(db.dispose)
+        await token_store.start()
+        stack.push_async_callback(token_store.stop)
+        internal_app, internal_urls, mcp_servers = _build_internal_mcp(
+            gateway.mcps, settings=settings
+        )
+        mcp_tools = await _prefetch_mcp_tools(mcp_servers)
+
+        pool = SessionPool()
+        bus = EventBus()
+        late = _LateConversations()
+        session_store = PostgresSessionStore(db)
+        backends = await _build_backends(
+            settings=settings,
+            agents=agents,
+            stack=stack,
+            db=db,
+            session_store=session_store,
+            mcp_internal_urls=internal_urls,
+            mcp_servers=mcp_servers,
+            mcp_tools=mcp_tools,
+            pool=pool,
+            late=late,
+        )
+        conversations = Conversations(
+            db=db,
+            agents=agents,
+            backends=backends,
+            bus=bus,
+            pool=pool,
+            store=session_store,
+            texts=gateway.texts,
+            frontends=gateway.frontends,
+            envelope=Envelope(
+                watch=gateway.watch, tz=gateway.tz, recall=gateway.recall
+            ),
+            distiller=gateway.distiller,
+            user_sink=gateway.user_sink,
+        )
+        late.conversations = conversations
+        scheduler = Scheduler(
+            conversations=conversations,
+            jobs=gateway.jobs,
+            driver=await _pgqueuer_driver(settings.database_url, stack),
+            budget=gateway.budget,
+            rotation=Rotation(
+                conversations, gateway.rotation or RotationPolicy(tz=gateway.tz)
+            ),
+            tz=gateway.tz,
+        )
+        conversations.scheduler = scheduler
+
+        runtime = GatewayRuntime(
+            agents=agents,
+            mcps=mcps,
+            backends=backends,
+            token_store=token_store,
+            db=db,
+            mcp_internal_urls=internal_urls,
+            admin_user=settings.admin_user,
+            admin_pass=settings.admin_pass,
+            session_secret=settings.session_secret,
+            frontends=tuple(gateway.frontends),
+            conversations=conversations,
+            bus=bus,
+            pool=pool,
+            scheduler=scheduler,
+            public_url=gateway.public_url.rstrip("/") if gateway.public_url else None,
+        )
+        for fe in gateway.frontends:
+            fe.configure(runtime)
+
+        _log.info(
+            "beaver-gateway: loaded %d agents, %d mcps, %d frontends",
+            len(agents),
+            len(mcps),
+            len(gateway.frontends),
+        )
+        if not gateway.frontends:
+            return
+
+        await conversations.start()
+        stack.push_async_callback(conversations.stop)
+        await scheduler.start()
+        stack.push_async_callback(scheduler.stop)
+        hooks = scheduler.app(
+            functools.partial(_authorize_hook, runtime=runtime, scope="api")
+        )
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(pool.reap_loop())
+            if internal_app is not None:
+                tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
+            tg.create_task(_serve_root(gateway, extra={"/hooks": hooks}))
+            if gateway.watch is not None:
+                tg.create_task(gateway.watch.run())
+            for fe in gateway.frontends:
+                tg.create_task(fe.serve())
+
+
+async def _authorize_hook(
+    request: Request, *, runtime: GatewayRuntime, scope: str
+) -> str:
+    return await require_token(request, runtime, scope=scope)
+
+
+async def _pgqueuer_driver(url: str, stack: AsyncExitStack) -> PsycopgDriver | None:
+    plain = _plain_postgres_url(url)
+    if plain is None:
+        return None
+    conn = await psycopg.AsyncConnection.connect(plain, autocommit=True)
+    stack.push_async_callback(conn.close)
+    return PsycopgDriver(conn)
+
+
+def _plain_postgres_url(url: str) -> str | None:
+    for prefix in ("postgresql+psycopg://", "postgresql://", "postgres://"):
+        if url.startswith(prefix):
+            return "postgresql://" + url[len(prefix) :]
+    return None
+
+
+async def _serve_root(gateway: Gateway, *, extra: dict[str, ASGIApp]) -> None:
+    app = build_root_app(gateway.frontends, extra=extra)
+    config = uvicorn.Config(app, host=gateway.host, port=gateway.port, log_level="info")
+    _log.info(
+        "gateway on http://%s:%d - %s",
+        gateway.host,
+        gateway.port,
+        ", ".join([*(fe.path for fe in gateway.frontends if fe.path), *extra])
+        or "no http frontends",
+    )
+    await uvicorn.Server(config).serve()
+
+
+def _build_internal_mcp(
+    mcps: list[McpServerT], *, settings: Settings
+) -> tuple[Starlette | None, dict[str, str], dict[str, FastMCP]]:
+    if not mcps:
+        return None, {}, {}
+    return build_internal_app(mcps, host="127.0.0.1", port=settings.internal_mcp_port)
+
+
+async def _prefetch_mcp_tools(
+    servers: dict[str, FastMCP],
+) -> dict[str, list[FastMCPTool]]:
+    out: dict[str, list[FastMCPTool]] = {}
+    for name, server in servers.items():
+        try:
+            out[name] = list(await server.list_tools())
+        except Exception:  # noqa: BLE001
+            _log.exception("failed to list tools for MCP %r, skipping", name)
+            out[name] = []
+    return out
+
+
+async def _serve_internal_mcp(app: Starlette, *, settings: Settings) -> None:
+    config = uvicorn.Config(
+        app,
+        host="127.0.0.1",
+        port=settings.internal_mcp_port,
+        log_level="warning",
+        loop="uvloop",
+    )
+    _log.info(
+        "internal MCP aggregator on http://127.0.0.1:%d/mcp/<name>",
+        settings.internal_mcp_port,
+    )
+    await uvicorn.Server(config).serve()
+
+
+class _LateConversations:
+    conversations: Conversations | None = None
+
+    def server(
+        self, key: str, _kind: str, names: tuple[str, ...]
+    ) -> McpSdkServerConfig | None:
+        if self.conversations is None or not names:
+            return None
+        return build_tool_server(self.conversations, conversation_key=key, names=names)
+
+    async def ask(self, key: str, payload: dict[str, Any]) -> str:
+        if self.conversations is None:
+            msg = "conversations service is not up yet"
+            raise RuntimeError(msg)
+        answer = await self.conversations.ask(key, payload)
+        return self.conversations.answer_text(answer)
+
+
+async def _build_backends(
+    *,
+    settings: Settings,
+    agents: AgentRegistry,
+    stack: AsyncExitStack,
+    db: Database,
+    session_store: PostgresSessionStore,
+    mcp_internal_urls: dict[str, str],
+    mcp_servers: dict[str, FastMCP],
+    mcp_tools: dict[str, list[FastMCPTool]],
+    pool: SessionPool,
+    late: _LateConversations,
+) -> dict[str, Backend]:
+    backends: dict[str, Backend] = {}
+
+    raycast_agents = [a for a in agents if isinstance(a, RaycastAgent)]
+    if raycast_agents:
+        client = await _try_open_raycast_client(settings, stack)
+        if client is not None:
+            raycast_backend = RaycastBackend(
+                client, mcp_servers=mcp_servers, mcp_tools=mcp_tools
+            )
+            for a in raycast_agents:
+                backends[a.name] = raycast_backend
+
+    runner = RunnerConfig(user=settings.claude_runner_user, home=settings.claude_home)
+    mcp_tool_names = {
+        name: [t.name for t in tools] for name, tools in mcp_tools.items()
+    }
+
+    async def record_usage(event: UsageEvent) -> None:
+        row = Usage(
+            agent_name=event.agent_name,
+            conversation_id=event.conversation_id,
+            session_id=event.session_id,
+            model=event.model,
+            effort=event.effort,
+            input_tokens=event.usage.input_tokens,
+            output_tokens=event.usage.output_tokens,
+            cache_read_tokens=event.usage.cache_read_tokens,
+            cache_creation_tokens=event.usage.cache_creation_tokens,
+            context_tokens=event.usage.context_tokens,
+            cost_usd=event.usage.cost_usd,
+            duration_ms=event.usage.duration_ms,
+            num_turns=event.usage.num_turns,
+            model_usage=event.usage.model_usage,
+        )
+        try:
+            async with db.session() as session:
+                await append_usage(session, row)
+        except Exception:  # noqa: BLE001
+            _log.exception("usage write failed for %s", event.agent_name)
+
+    async def record_tool(event: ToolAudit) -> None:
+        detail = {
+            "conversation": event.conversation,
+            "kind": event.kind,
+            "tool": event.tool,
+            "decision": event.decision,
+            "reason": event.reason,
+            "brief": event.brief,
+        }
+        try:
+            async with db.session() as session:
+                await append_audit(
+                    session,
+                    actor=f"agent:{event.agent}",
+                    kind="tool_call",
+                    agent_name=event.agent,
+                    detail=detail,
+                )
+        except Exception:  # noqa: BLE001
+            _log.exception("tool audit write failed for %s", event.agent)
+
+    for a in agents:
+        if isinstance(a, ClaudeAgent):
+            adapter = ClaudeSdkBackend(
+                agent=a,
+                mcp_internal_urls=mcp_internal_urls,
+                session_store=session_store,
+                mcp_tool_names=mcp_tool_names,
+                runner=runner,
+                usage_sink=record_usage,
+                pool=pool,
+                tool_server=functools.partial(late.server, names=a.gateway_tools),
+                asker=late.ask,
+                audit_sink=record_tool,
+            )
+            await stack.enter_async_context(adapter)
+            backends[a.name] = adapter
+
+    return backends
+
+
+async def _try_open_raycast_client(
+    settings: Settings, stack: AsyncExitStack
+) -> RaycastClient | None:
+    if not settings.raycast_bearer:
+        _log.warning(
+            "RaycastAgent present but RAYCAST_BEARER is unset; those agents 503"
+        )
+        return None
+    if not settings.raycast_device_id:
+        _log.warning(
+            "RaycastAgent present but RAYCAST_DEVICE_ID is unset, those agents 503 "
+            "(generate once with `python -c 'import secrets; "
+            "print(secrets.token_hex(32))'`)"
+        )
+        return None
+    if not settings.raycast_config_path.exists():
+        _log.warning(
+            "RaycastAgent present but %s is missing, those agents 503",
+            settings.raycast_config_path,
+        )
+        return None
+
+    config = RaycastConfig.load(settings.raycast_config_path)
+    client = RaycastClient(
+        config=config,
+        bearer_token=settings.raycast_bearer,
+        device_id=settings.raycast_device_id,
+        locale=settings.raycast_locale,
+    )
+    return await stack.enter_async_context(client)