bitrouter_runtime/
server.rs

1use std::sync::Arc;
2
3use bitrouter_api::router::{anthropic, google, 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        let generate_content = auth_gate(auth::openai_auth(auth_ctx.clone())).and(
86            google::generate_content::filters::generate_content_filter(
87                self.table.clone(),
88                self.router.clone(),
89            ),
90        );
91
92        // Key management routes — always mounted (returns 404 if no DB, since
93        // the filter will not match without the DB anyway).
94        let key_mgmt = keys::key_routes(auth_ctx.clone(), self.db.clone());
95
96        let routes = health
97            .or(chat)
98            .or(messages)
99            .or(responses)
100            .or(generate_content)
101            .or(key_mgmt)
102            .recover(handle_auth_rejection)
103            .with(warp::trace::request());
104
105        let server = warp::serve(routes)
106            .bind(addr)
107            .await
108            .graceful(shutdown_signal());
109
110        if auth_ctx.is_open() {
111            tracing::info!(%addr, "server listening (auth disabled — no master_key configured)");
112        } else {
113            tracing::info!(%addr, "server listening (auth enabled)");
114        }
115        server.run().await;
116        tracing::info!("server stopped");
117
118        Ok(())
119    }
120}
121
122/// Convert an auth filter into a gate that rejects unauthorized requests
123/// but does not add anything to the extract tuple. This lets us compose
124/// `auth_gate(auth).and(existing_filter)` without changing the existing
125/// filter's handler signature.
126fn auth_gate(
127    auth: impl Filter<Extract = (bitrouter_accounts::identity::Identity,), Error = warp::Rejection>
128    + Clone,
129) -> impl Filter<Extract = (), Error = warp::Rejection> + Clone {
130    auth.map(|_| ()).untuple_one()
131}
132
133/// Rejection handler that turns [`Unauthorized`] into a JSON 401 response.
134async fn handle_auth_rejection(
135    rejection: warp::Rejection,
136) -> std::result::Result<impl warp::Reply, warp::Rejection> {
137    if let Some(e) = rejection.find::<Unauthorized>() {
138        let json = warp::reply::json(&serde_json::json!({
139            "error": {
140                "message": e.to_string(),
141                "type": "authentication_error",
142            }
143        }));
144        return Ok(warp::reply::with_status(
145            json,
146            warp::http::StatusCode::UNAUTHORIZED,
147        ));
148    }
149    Err(rejection)
150}
151
152async fn shutdown_signal() {
153    let ctrl_c = tokio::signal::ctrl_c();
154
155    #[cfg(unix)]
156    {
157        let mut term =
158            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()).unwrap();
159        tokio::select! {
160            _ = ctrl_c => {}
161            _ = term.recv() => {}
162        }
163    }
164
165    #[cfg(not(unix))]
166    {
167        ctrl_c.await.ok();
168    }
169
170    tracing::info!("shutdown signal received");
171}
bitrouter_runtime/server.rs

bitrouter_runtime/
server.rs
