442 lines
15 KiB
Python
442 lines
15 KiB
Python
"""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, scopes_with
|
|
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,
|
|
scopes=scopes_with(fe.scope for fe in gateway.frontends),
|
|
)
|
|
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)
|