diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index ba7ebcd8008..f170f9b290e 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1679,6 +1679,7 @@ mod tests { use super::*; use crate::auth::JwtAuthProvider; use crate::routes::subscribe::{HasWebSocketOptions, WebSocketOptions}; + use crate::routes::{identity::IdentityRoutes, router_with_root_routes, RootRoutes}; use crate::{ Action, Authorization, ControlStateReadAccess, ControlStateWriteAccess, MaybeMisdirected, Unauthorized, }; @@ -2452,4 +2453,79 @@ mod tests { remove_http_response_size_metric(database_identity); } + + fn root_router(root_routes: RootRoutes) -> axum::Router { + let state = DummyState::new(); + router_with_root_routes( + &state, + DatabaseRoutes::default(), + IdentityRoutes::default(), + root_routes, + axum::Router::new(), + ) + .with_state(state) + } + + fn post_mcp_root(body: &'static str) -> Request { + Request::builder() + .method(http::Method::POST) + .uri("/v1/mcp") + .header(http::header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap() + } + + #[tokio::test] + async fn default_root_routes_serve_the_real_handlers() { + let app = root_router(RootRoutes::default()); + + let response = app + .clone() + .oneshot(post_mcp_root( + r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"ping"}}"#, + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response.into_body().collect().await.unwrap().to_bytes(); + assert!(std::str::from_utf8(&body).unwrap().contains("pong")); + + let response = app + .oneshot(Request::builder().uri("/v1/ping").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn a_substituted_root_mcp_route_replaces_the_default_handler() { + let app = root_router(RootRoutes { + mcp_post: axum::routing::post(|| async { "substituted" }), + ..Default::default() + }); + + let response = app.oneshot(post_mcp_root("")).await.unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.into_body().collect().await.unwrap().to_bytes(), "substituted"); + } + + #[tokio::test] + async fn the_auth_middleware_runs_before_substituted_root_layers() { + let app = root_router(RootRoutes { + mcp_post: axum::routing::post(|| async { "substituted" }).layer(axum::middleware::from_fn( + |request: axum::extract::Request, next: axum::middleware::Next| async move { + if request.extensions().get::().is_none() { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + next.run(request).await + }, + )), + ..Default::default() + }); + + let response = app.oneshot(post_mcp_root("")).await.unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + } } diff --git a/crates/client-api/src/routes/mod.rs b/crates/client-api/src/routes/mod.rs index 19b92e19402..5c487aa5fd9 100644 --- a/crates/client-api/src/routes/mod.rs +++ b/crates/client-api/src/routes/mod.rs @@ -1,3 +1,4 @@ +use axum::routing::MethodRouter; use http::header; use tower_http::cors; @@ -20,7 +21,27 @@ use self::{database::DatabaseRoutes, identity::IdentityRoutes}; /// establish a connection to SpacetimeDB. This API call doesn't actually do anything. pub async fn ping(_auth: crate::auth::SpacetimeAuthHeader) {} -#[allow(clippy::let_and_return)] +/// Allows the edition to customize the routes directly under `/v1`, as [`DatabaseRoutes`] does for `/database`. +pub struct RootRoutes { + /// GET: /ping + pub ping_get: MethodRouter, + /// POST: /mcp + pub mcp_post: MethodRouter, +} + +impl Default for RootRoutes +where + S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, +{ + fn default() -> Self { + use axum::routing::{get, post}; + Self { + ping_get: get(ping), + mcp_post: post(mcp::mcp_root::), + } + } +} + pub fn router( ctx: &S, database_routes: DatabaseRoutes, @@ -30,7 +51,19 @@ pub fn router( where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { - use axum::routing::{get, post}; + router_with_root_routes(ctx, database_routes, identity_routes, RootRoutes::default(), extra) +} + +pub fn router_with_root_routes( + ctx: &S, + database_routes: DatabaseRoutes, + identity_routes: IdentityRoutes, + root_routes: RootRoutes, + extra: axum::Router, +) -> axum::Router +where + S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, +{ let router = axum::Router::new() .nest("/database", database_routes.into_router(ctx.clone())) .nest("/identity", identity_routes.into_router()) @@ -40,12 +73,12 @@ where // the database is named in the request body, so `mcp_root` counts its own egress .route( "/mcp", - post(mcp::mcp_root::).route_layer(axum::middleware::from_fn_with_state( + root_routes.mcp_post.route_layer(axum::middleware::from_fn_with_state( ctx.clone(), crate::auth::anon_auth_middleware::, )), ) - .route("/ping", get(ping)) + .route("/ping", root_routes.ping_get) .merge(extra); let cors = cors::CorsLayer::new()