369 lines
14 KiB
Python
369 lines
14 KiB
Python
"""Process entrypoint.
|
|
|
|
Phase 1.4 — async ``main``: install uvloop, load the user config, build
|
|
registries + per-agent backends (only ``RaycastBackend`` so far), wire
|
|
each frontend with a ``GatewayRuntime``, and run all
|
|
``frontend.serve()`` coroutines concurrently. Without any frontends we
|
|
still print the Phase 0 DoD line and exit cleanly so the bare skeleton
|
|
keeps working.
|
|
|
|
Phase 2.1 — when the user declares any ``McpServer``, we additionally
|
|
build the internal MCP aggregator app and run it on
|
|
``127.0.0.1:INTERNAL_MCP_PORT`` as another task inside the same
|
|
TaskGroup. URLs are surfaced through ``GatewayRuntime.mcp_internal_urls``
|
|
so Phase 2.2's ClaudeCode adapter can find them.
|
|
|
|
Phase 3 — the same aggregator backs the external ``McpServerFrontend``;
|
|
``cli`` doesn't have to know that the frontend reverse-proxies into it,
|
|
it just keeps the aggregator running for anyone who needs it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import signal
|
|
from contextlib import AsyncExitStack
|
|
from typing import TYPE_CHECKING
|
|
|
|
import uvicorn
|
|
import uvloop
|
|
from dotenv import load_dotenv
|
|
from raycast_api import Client as RaycastClient
|
|
from raycast_api.config import Config as RaycastConfig
|
|
|
|
from beaver_gateway import config_loader
|
|
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.core.auth import TokenStore
|
|
from beaver_gateway.core.registry import AgentRegistry, McpRegistry
|
|
from beaver_gateway.frontends.base import GatewayRuntime
|
|
from beaver_gateway.mcp.internal_app import build_internal_app
|
|
from beaver_gateway.settings import Settings
|
|
from beaver_gateway.storage import Database, PostgresSessionStore, Usage, append_usage
|
|
|
|
if TYPE_CHECKING:
|
|
from fastmcp import FastMCP
|
|
from fastmcp.tools.base import Tool as FastMCPTool
|
|
from starlette.applications import Starlette
|
|
|
|
from beaver_gateway.backends.base import Backend
|
|
from beaver_gateway.mcp.types import McpServerT
|
|
|
|
|
|
_log = logging.getLogger("beaver_gateway.cli")
|
|
|
|
|
|
def main() -> None:
|
|
"""Sync wrapper: uvloop loop factory + asyncio.run."""
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
|
|
)
|
|
_install_sigterm_handler()
|
|
asyncio.run(_async_main(), loop_factory=uvloop.new_event_loop)
|
|
|
|
|
|
def _install_sigterm_handler() -> None:
|
|
"""Turn SIGTERM into a normal interpreter exit.
|
|
|
|
Python's default SIGTERM disposition kills the process outright, so
|
|
neither ``AsyncExitStack`` unwinding nor ``atexit`` hooks run. That
|
|
matters because every ``claude`` we spawn lives in its own session
|
|
(ptyprocess calls ``setsid``), which makes it immune to the signal
|
|
that took us down — a hard SIGTERM leaves one orphaned CLI per live
|
|
session, each holding hundreds of MB. Raising ``KeyboardInterrupt``
|
|
instead routes ``docker stop`` / ``systemctl stop`` through the same
|
|
shutdown path as Ctrl-C, which does reap them.
|
|
|
|
Only installed when we own the main thread's signal handlers; under
|
|
an embedding host that isn't ours to take.
|
|
"""
|
|
|
|
def _raise_interrupt(_signum: int, _frame: object) -> None:
|
|
raise KeyboardInterrupt
|
|
|
|
with contextlib.suppress(ValueError, OSError):
|
|
signal.signal(signal.SIGTERM, _raise_interrupt)
|
|
|
|
|
|
async def _async_main() -> None:
|
|
# Populate ``os.environ`` from ``.env`` before anything else so the
|
|
# user's ``config.py`` can read its own secrets via ``os.environ[...]``
|
|
# (Firefly PAT, third-party MCP creds, etc.). ``Settings`` already
|
|
# reads ``.env`` independently via pydantic-settings, but that path
|
|
# populates Settings fields, not the process environment.
|
|
# ``override=False``: real env vars (Docker, systemd) win over .env.
|
|
load_dotenv(override=False)
|
|
|
|
settings = Settings() # ty: ignore[missing-argument]
|
|
gateway = config_loader.load(settings.config_path)
|
|
|
|
agents = AgentRegistry(gateway.agents)
|
|
mcps = McpRegistry(gateway.mcps)
|
|
|
|
# Phase 4.1 — open the async DB and run create_all once. Engine
|
|
# pool is process-wide; ``dispose()`` after the TaskGroup unwinds.
|
|
db = Database(settings.database_url)
|
|
await db.create_all()
|
|
|
|
# Phase 4.2 — TokenStore now reads from the DB (in-memory cache
|
|
# primed at start, TTL-refreshed, last_used_at flushed by a
|
|
# background task). BOOTSTRAP_TOKENS layers on top so first-run /
|
|
# examples still work without DB writes.
|
|
token_store = TokenStore(
|
|
db, bootstrap=TokenStore.parse_bootstrap(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 MCP URLs must exist before we construct any
|
|
# ClaudeSdkBackend - adapters bake the URLs into their
|
|
# ``mcp_servers`` at construction time. The
|
|
# ``mcp_servers`` map is used by the Raycast backend, which
|
|
# needs in-process ``list_tools`` / ``call_tool`` access (the
|
|
# Raycast wire has no native MCP concept).
|
|
internal_app, internal_urls, mcp_servers = _build_internal_mcp(
|
|
gateway.mcps, settings=settings
|
|
)
|
|
|
|
# Prefetch tool catalogs for every MCP so RaycastAgent requests
|
|
# don't pay a per-turn list_tools roundtrip and so a broken MCP
|
|
# surfaces at startup instead of mid-conversation.
|
|
mcp_tools = await _prefetch_mcp_tools(mcp_servers)
|
|
|
|
backends: dict[str, Backend] = await _build_backends(
|
|
settings=settings,
|
|
agents=agents,
|
|
stack=stack,
|
|
db=db,
|
|
mcp_internal_urls=internal_urls,
|
|
mcp_servers=mcp_servers,
|
|
mcp_tools=mcp_tools,
|
|
)
|
|
|
|
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),
|
|
)
|
|
|
|
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),
|
|
)
|
|
|
|
# Keep the Phase 0 DoD line on stdout for grep-friendly smoke
|
|
# tests, in addition to the structured log line above.
|
|
print(
|
|
f"beaver-gateway: loaded {len(agents)} agents, "
|
|
f"{len(mcps)} mcps, {len(gateway.frontends)} frontends"
|
|
)
|
|
|
|
if not gateway.frontends:
|
|
# No external listeners → nothing to serve. The internal
|
|
# MCP app has no consumer on its own, so we skip running
|
|
# it in this path and exit cleanly (Phase 0 DoD).
|
|
return
|
|
|
|
async with asyncio.TaskGroup() as tg:
|
|
if internal_app is not None:
|
|
tg.create_task(_serve_internal_mcp(internal_app, settings=settings))
|
|
for fe in gateway.frontends:
|
|
tg.create_task(fe.serve())
|
|
|
|
|
|
def _build_internal_mcp(
|
|
mcps: list[McpServerT], *, settings: Settings
|
|
) -> tuple[Starlette | None, dict[str, str], dict[str, FastMCP]]:
|
|
"""Build the aggregator app + URL map + server map, or empty equivalents.
|
|
|
|
The URL map is always handed out (frontends may still introspect
|
|
``runtime.mcp_internal_urls`` even if nothing is configured); the
|
|
app is ``None`` when there are no MCPs to mount, so the caller
|
|
skips the uvicorn task entirely. The server map is the in-process
|
|
handle the Raycast backend needs to splice MCP tools into its
|
|
requests — empty when no MCPs are configured.
|
|
"""
|
|
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]]:
|
|
"""Eagerly enumerate tools per MCP so the Raycast loop has a static catalog.
|
|
|
|
Each underlying proxy is allowed to fail independently — a broken
|
|
MCP shouldn't take down the whole gateway. The result has one entry
|
|
per MCP that responded; agents that ``expose_mcps`` a missing entry
|
|
will simply expose no tools from it (logged once per request).
|
|
"""
|
|
out: dict[str, list[FastMCPTool]] = {}
|
|
for name, server in servers.items():
|
|
try:
|
|
out[name] = list(await server.list_tools())
|
|
except Exception: # noqa: BLE001 — proxy can raise any transport error; we degrade per-MCP rather than fail the whole gateway
|
|
_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:
|
|
"""Run the internal MCP aggregator on loopback.
|
|
|
|
Bound to ``127.0.0.1`` (never EXPOSE'd) — only the in-process
|
|
ClaudeCode subprocess reaches it. Logged at ``warning`` level so
|
|
we don't drown the gateway's own logs in per-request noise.
|
|
"""
|
|
config = uvicorn.Config(
|
|
app,
|
|
host="127.0.0.1",
|
|
port=settings.internal_mcp_port,
|
|
log_level="warning",
|
|
loop="uvloop",
|
|
)
|
|
server = uvicorn.Server(config)
|
|
_log.info(
|
|
"internal MCP aggregator on http://127.0.0.1:%d/mcp/<name>",
|
|
settings.internal_mcp_port,
|
|
)
|
|
await server.serve()
|
|
|
|
|
|
async def _build_backends(
|
|
*,
|
|
settings: Settings,
|
|
agents: AgentRegistry,
|
|
stack: AsyncExitStack,
|
|
db: Database,
|
|
mcp_internal_urls: dict[str, str],
|
|
mcp_servers: dict[str, FastMCP],
|
|
mcp_tools: dict[str, list[FastMCPTool]],
|
|
) -> dict[str, Backend]:
|
|
"""Construct one backend per agent name.
|
|
|
|
The Raycast ``Client`` is shared across every ``RaycastAgent``
|
|
(bearer + device-id are process-wide), so we open it lazily - only
|
|
when at least one ``RaycastAgent`` is present - and close it via
|
|
the caller's exit stack.
|
|
|
|
Each :class:`ClaudeAgent` gets its own :class:`ClaudeSdkBackend`
|
|
(own live-session pool, own prompt and MCP set); all of them share
|
|
the session store and the usage sink.
|
|
"""
|
|
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
|
|
|
|
session_store = PostgresSessionStore(db)
|
|
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,
|
|
cost_usd=event.usage.cost_usd,
|
|
duration_ms=event.usage.duration_ms,
|
|
num_turns=event.usage.num_turns,
|
|
)
|
|
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)
|
|
|
|
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,
|
|
)
|
|
await stack.enter_async_context(adapter)
|
|
backends[a.name] = adapter
|
|
|
|
return backends
|
|
|
|
|
|
async def _try_open_raycast_client(
|
|
settings: Settings, stack: AsyncExitStack
|
|
) -> RaycastClient | None:
|
|
"""Open a ``raycast_api.Client`` if creds are available, else warn + skip.
|
|
|
|
Skipping is gentler than failing startup: the user can bring up the
|
|
gateway, list their agents, and still hit a ``ClaudeAgent``; affected
|
|
``RaycastAgent`` instances 503 with a specific message at request time.
|
|
Missing ``raycast.json`` falls into the same bucket — first-run
|
|
users won't have it yet.
|
|
"""
|
|
if not settings.raycast_bearer:
|
|
_log.warning(
|
|
"RaycastAgent present but RAYCAST_BEARER is unset — those agents will 503"
|
|
)
|
|
return None
|
|
if not settings.raycast_device_id:
|
|
_log.warning(
|
|
"RaycastAgent present but RAYCAST_DEVICE_ID is unset — those agents "
|
|
"will 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 will 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)
|