bitrouter_runtime/
server.rs

1use std::sync::Arc;
2
3use bitrouter_api::router::{anthropic, openai};
4use bitrouter_config::BitrouterConfig;
5use bitrouter_core::routers::{model_router::LanguageModelRouter, routing_table::RoutingTable};
6use sea_orm::DatabaseConnection;
7use warp::Filter;
8
9use crate::auth::{self, AuthContext, Unauthorized};
10use crate::error::Result;
11use crate::keys;
12
13/// A stub model router that rejects all requests with a descriptive error.
14///
15/// Used when the server starts without a real provider-backed router. Health
16/// checks and other non-model endpoints still work; only model API requests
17/// will return an error.
18pub struct StubModelRouter;
19
20impl LanguageModelRouter for StubModelRouter {
21    async fn route_model(
22        &self,
23        _target: bitrouter_core::routers::routing_table::RoutingTarget,
24    ) -> bitrouter_core::errors::Result<
25        Box<bitrouter_core::models::language::language_model::DynLanguageModel<'static>>,
26    > {
27        Err(bitrouter_core::errors::BitrouterError::unsupported(
28            "runtime",
29            "model routing",
30            Some("no model router configured — configure providers to enable API endpoints".into()),
31        ))
32    }
33}
34
35pub struct ServerPlan<T, R> {
36    config: BitrouterConfig,
37    table: Arc<T>,
38    router: Arc<R>,
39    db: Option<Arc<DatabaseConnection>>,
40}
41
42impl<T, R> ServerPlan<T, R>
43where
44    T: RoutingTable + Send + Sync + 'static,
45    R: LanguageModelRouter + Send + Sync + 'static,
46{
47    pub fn new(config: BitrouterConfig, table: Arc<T>, router: Arc<R>) -> Self {
48        Self {
49            config,
50            table,
51            router,
52            db: None,
53        }
54    }
55
56    /// Set the database connection for virtual key lookups and key management.
57    pub fn with_db(mut self, db: DatabaseConnection) -> Self {
58        self.db = Some(Arc::new(db));
59        self
60    }
61
62    pub async fn serve(self) -> Result<()> {
63        let addr = self.config.server.listen;
64
65        // Build auth context.
66        let auth_ctx = Arc::new(AuthContext::new(
67            self.config.master_key.as_deref(),
68            self.db.as_ref().map(|db| db.as_ref().clone()),
69        ));
70
71        let health = warp::path("health")
72            .and(warp::get())
73            .map(|| warp::reply::json(&serde_json::json!({ "status": "ok" })));
74
75        // Model API routes — gated by protocol-appropriate auth.
76        let chat = auth_gate(auth::openai_auth(auth_ctx.clone())).and(
77            openai::chat::filters::chat_completions_filter(self.table.clone(), self.router.clone()),
78        );
79        let messages = auth_gate(auth::anthropic_auth(auth_ctx.clone())).and(
80            anthropic::messages::filters::messages_filter(self.table.clone(), self.router.clone()),
81        );
82        let responses = auth_gate(auth::openai_auth(auth_ctx.clone())).and(
83            openai::responses::filters::responses_filter(self.table.clone(), self.router.clone()),
84        );
85
86        // Key management routes — always mounted (returns 404 if no DB, since
87        // the filter will not match without the DB anyway).
88        let key_mgmt = keys::key_routes(auth_ctx.clone(), self.db.clone());
89
90        let routes = health
91            .or(chat)
92            .or(messages)
93            .or(responses)
94            .or(key_mgmt)
95            .recover(handle_auth_rejection)
96            .with(warp::trace::request());
97
98        let server = warp::serve(routes)
99            .bind(addr)
100            .await
101            .graceful(shutdown_signal());
102
103        if auth_ctx.is_open() {
104            tracing::info!(%addr, "server listening (auth disabled — no master_key configured)");
105        } else {
106            tracing::info!(%addr, "server listening (auth enabled)");
107        }
108        server.run().await;
109        tracing::info!("server stopped");
110
111        Ok(())
112    }
113}
114
115/// Convert an auth filter into a gate that rejects unauthorized requests
116/// but does not add anything to the extract tuple. This lets us compose
117/// `auth_gate(auth).and(existing_filter)` without changing the existing
118/// filter's handler signature.
119fn auth_gate(
120    auth: impl Filter<Extract = (bitrouter_accounts::identity::Identity,), Error = warp::Rejection>
121    + Clone,
122) -> impl Filter<Extract = (), Error = warp::Rejection> + Clone {
123    auth.map(|_| ()).untuple_one()
124}
125
126/// Rejection handler that turns [`Unauthorized`] into a JSON 401 response.
127async fn handle_auth_rejection(
128    rejection: warp::Rejection,
129) -> std::result::Result<impl warp::Reply, warp::Rejection> {
130    if let Some(e) = rejection.find::<Unauthorized>() {
131        let json = warp::reply::json(&serde_json::json!({
132            "error": {
133                "message": e.to_string(),
134                "type": "authentication_error",
135            }
136        }));
137        return Ok(warp::reply::with_status(
138            json,
139            warp::http::StatusCode::UNAUTHORIZED,
140        ));
141    }
142    Err(rejection)
143}
144
145async fn shutdown_signal() {
146    let ctrl_c = tokio::signal::ctrl_c();
147
148    #[cfg(unix)]
149    {
150        let mut term =
151            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
152        tokio::select! {
153            _ = ctrl_c => {}
154            _ = term.recv() => {}
155        }
156    }
157
158    #[cfg(not(unix))]
159    {
160        ctrl_c.await.ok();
161    }
162
163    tracing::info!("shutdown signal received");
164}
bitrouter_runtime/server.rs

bitrouter_runtime/
server.rs
