ai_crew_sync/
serve.rs

1use std::time::Duration;
2
3use anyhow::Context;
4use axum::response::IntoResponse;
5use axum::{
6    Router,
7    extract::State,
8    routing::{get, post},
9};
10use rmcp::transport::streamable_http_server::{
11    StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
12};
13use sqlx::PgPool;
14use tokio_util::sync::CancellationToken;
15use tower_http::{limit::RequestBodyLimitLayer, trace::TraceLayer};
16
17use crate::{auth, dashboard, events::EventHub, tools::Bus, webhooks};
18
19pub struct ServeOptions {
20    pub bind: String,
21    /// Hostnames accepted in the `Host` header. Empty disables the check, which
22    /// is what you want behind a proxy that already validates it.
23    pub allowed_hosts: Vec<String>,
24    pub allowed_origins: Vec<String>,
25    /// Largest MCP request body accepted, in bytes.
26    pub max_request_bytes: usize,
27    /// Requests per minute per token; 0 disables in-process limiting.
28    pub rate_limit_per_minute: u32,
29    /// Signs the dashboard's read-only session grants. Share it across
30    /// replicas so a grant issued by one is accepted by the others.
31    pub dashboard_secret: Vec<u8>,
32}
33
34/// Headroom over the largest legitimate request: a 1 MiB message body plus
35/// eight 256 KiB attachments, which base64 inflates to ~2.7 MiB, plus JSON
36/// framing. Anything beyond this is a mistake or an attack, and rejecting it
37/// before parsing keeps a bad request from costing a large allocation.
38pub const DEFAULT_MAX_REQUEST_BYTES: usize = 8 * 1024 * 1024;
39pub const DEFAULT_RATE_LIMIT_PER_MINUTE: u32 = 600;
40
41async fn health(
42    State(pool): State<PgPool>,
43) -> (axum::http::StatusCode, axum::Json<serde_json::Value>) {
44    match sqlx::query_scalar::<_, i32>("SELECT 1")
45        .fetch_one(&pool)
46        .await
47    {
48        Ok(_) => (
49            axum::http::StatusCode::OK,
50            axum::Json(serde_json::json!({ "status": "ok", "database": "up" })),
51        ),
52        Err(e) => {
53            tracing::error!(error = %e, "health check failed");
54            (
55                axum::http::StatusCode::SERVICE_UNAVAILABLE,
56                axum::Json(serde_json::json!({ "status": "degraded", "database": "down" })),
57            )
58        }
59    }
60}
61
62/// `DefaultBodyLimit` answers with a bare 413. The caller here is a language
63/// model, so replace it with a body that says what the limit is and what to
64/// do instead.
65async fn explain_payload_too_large(
66    limit: usize,
67    req: axum::extract::Request,
68    next: axum::middleware::Next,
69) -> axum::response::Response {
70    let resp = next.run(req).await;
71    if resp.status() != axum::http::StatusCode::PAYLOAD_TOO_LARGE {
72        return resp;
73    }
74    (
75        axum::http::StatusCode::PAYLOAD_TOO_LARGE,
76        axum::Json(serde_json::json!({
77            "error": format!(
78                "request body is too large; this server accepts up to {limit} bytes. \
79                 Send fewer or smaller attachments (256 KiB each, 8 per message), \
80                 or split the call into several smaller ones."
81            )
82        })),
83    )
84        .into_response()
85}
86
87pub fn build_router(pool: PgPool, opts: &ServeOptions, ct: CancellationToken) -> Router {
88    // One LISTEN connection feeds every in-process consumer: wait_for_updates
89    // long-polls and the webhook dispatcher.
90    let hub = EventHub::new();
91    tokio::spawn(crate::events::run_pg_listener(
92        pool.clone(),
93        hub.clone(),
94        ct.clone(),
95    ));
96    tokio::spawn(webhooks::run_dispatcher(
97        pool.clone(),
98        hub.clone(),
99        ct.clone(),
100    ));
101
102    let mut config = StreamableHttpServerConfig::default()
103        .with_json_response(true)
104        // Sessions are gone in the current MCP protocol revision; running
105        // stateless means any instance can serve any request, so the service
106        // scales horizontally behind a plain load balancer.
107        .with_legacy_session_mode(false)
108        .with_sse_keep_alive(Some(Duration::from_secs(30)))
109        .with_cancellation_token(ct);
110
111    config = if opts.allowed_hosts.is_empty() {
112        config.disable_allowed_hosts()
113    } else {
114        config.with_allowed_hosts(opts.allowed_hosts.clone())
115    };
116    if !opts.allowed_origins.is_empty() {
117        config = config.with_allowed_origins(opts.allowed_origins.clone());
118    }
119
120    let mcp: StreamableHttpService<Bus, LocalSessionManager> = StreamableHttpService::new(
121        {
122            let pool = pool.clone();
123            let hub = hub.clone();
124            move || Ok(Bus::new(pool.clone(), hub.clone()))
125        },
126        Default::default(),
127        config,
128    );
129
130    let limiter = crate::ratelimit::RateLimiter::new(opts.rate_limit_per_minute);
131    if limiter.is_none() {
132        tracing::warn!("in-process rate limiting is disabled (BUS_RATE_LIMIT_PER_MINUTE=0)");
133    }
134    let auth_state = auth::AuthState {
135        pool: pool.clone(),
136        limiter,
137    };
138
139    let mcp_routes = Router::new()
140        .nest_service("/mcp", mcp)
141        // Every /mcp request must carry a valid bearer token; the middleware
142        // injects the resolved AuthCtx that tool handlers read, and charges
143        // the per-token rate limit.
144        .layer(axum::middleware::from_fn_with_state(
145            auth_state,
146            auth::require_bearer,
147        ))
148        // Runs before auth: the body is truncated at the limit rather than
149        // read whole, and no token lookup happens. This must be tower-http's
150        // layer, not axum's DefaultBodyLimit — the MCP service reads the raw
151        // body itself, so an extractor-level limit would never fire.
152        .layer(RequestBodyLimitLayer::new(opts.max_request_bytes))
153        .layer(axum::middleware::from_fn({
154            let limit = opts.max_request_bytes;
155            move |req, next| explain_payload_too_large(limit, req, next)
156        }));
157
158    let dashboard_state = dashboard::DashboardState {
159        pool: pool.clone(),
160        secret: std::sync::Arc::new(opts.dashboard_secret.clone()),
161    };
162    let dashboard_routes = Router::new()
163        .route("/dashboard", get(dashboard::render))
164        .route("/dashboard/login", post(dashboard::login))
165        .with_state(dashboard_state);
166
167    Router::new()
168        .route("/health", get(health))
169        .merge(dashboard_routes)
170        .merge(mcp_routes)
171        // Record the method and path only. The default span records the whole
172        // URI, which is how a credential in a query string reaches the logs —
173        // the dashboard no longer accepts one, and the logs no longer invite it.
174        .layer(
175            TraceLayer::new_for_http().make_span_with(|req: &axum::http::Request<_>| {
176                tracing::info_span!(
177                    "http",
178                    method = %req.method(),
179                    path = %req.uri().path(),
180                )
181            }),
182        )
183        .with_state(pool)
184}
185
186pub async fn run(pool: PgPool, opts: ServeOptions) -> anyhow::Result<()> {
187    let ct = CancellationToken::new();
188    let app = build_router(pool, &opts, ct.child_token());
189
190    let listener = tokio::net::TcpListener::bind(&opts.bind)
191        .await
192        .with_context(|| format!("failed to bind {}", opts.bind))?;
193    let addr = listener.local_addr()?;
194    tracing::info!(%addr, "ai-crew-sync listening; MCP endpoint at /mcp");
195
196    let shutdown = {
197        let ct = ct.clone();
198        async move {
199            let ctrl_c = async {
200                tokio::signal::ctrl_c().await.ok();
201            };
202            #[cfg(unix)]
203            let term = async {
204                if let Ok(mut s) =
205                    tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
206                {
207                    s.recv().await;
208                }
209            };
210            #[cfg(not(unix))]
211            let term = std::future::pending::<()>();
212
213            tokio::select! {
214                _ = ctrl_c => {},
215                _ = term => {},
216            }
217            tracing::info!("shutdown signal received");
218            ct.cancel();
219        }
220    };
221
222    axum::serve(listener, app)
223        .with_graceful_shutdown(shutdown)
224        .await
225        .context("server error")?;
226    Ok(())
227}
ai_crew_sync/serve.rs

ai_crew_sync/
serve.rs
