From 0d24b961acde63c6c094684b44af555accaa674c Mon Sep 17 00:00:00 2001 From: Tien Pham Date: Fri, 21 Aug 2026 15:22:57 +0300 Subject: [PATCH] feat: make egress-tracking middleware apply to MCP route --- crates/client-api/src/routes/database.rs | 70 ++++++++++++++++-- crates/client-api/src/routes/mcp.rs | 90 +++++++++++++++--------- 2 files changed, 123 insertions(+), 37 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 6c1909949fe..fffc3ac0b9c 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1551,6 +1551,7 @@ where .route("/schema", self.schema_get) .route("/logs", self.logs_get) .route("/sql", self.sql_post) + .route("/mcp", self.mcp_post) .route("/unstable/timestamp", self.timestamp_get) .route("/pre_publish", self.pre_publish) .route("/reset", self.db_reset) @@ -1568,10 +1569,7 @@ where // so we don't mind that we don't measure them. let db_router = db_router .route("/", self.db_delete) - .route("/identity", self.identity_get) - // I (pgoldman 2026-08-07) am actually somewhat concerned that we do care about measuring egress for MCP requests, - // but the MCP handler's name resolution and error handling are significantly incompatible with the middleware. - .route("/mcp", self.mcp_post); + .route("/identity", self.identity_get); // Add the subscribe route after `resolving_egress_metrics_middleware` // so that its egress bytes don't get counted into `http_response_size_bytes`; @@ -2325,6 +2323,70 @@ mod tests { remove_http_response_size_metric(database_identity); } + #[tokio::test] + async fn http_response_egress_metric_counts_mcp() { + let database_identity = test_identity(20); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + mcp_post: axum::routing::post(|| async { + ([(http::header::CONTENT_TYPE, "application/json")], r#"{"result":{}}"#) + }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .method(http::Method::POST) + .uri(format!("/{database_identity}/mcp")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, r#"{"result":{}}"#); + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "application/json".len() as u64 + r#"{"result":{}}"#.len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn mcp_handshake_returns_not_found_for_an_unknown_database() { + let state = DummyState::new(); + let app = DatabaseRoutes:: { + mcp_post: axum::routing::post(|| async { "not reached" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .method(http::Method::POST) + .uri("/unregistered-name/mcp") + .body(Body::from(r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::NOT_FOUND); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "`unregistered-name` not found" + ); + } + #[tokio::test] async fn http_response_egress_metric_counts_error_responses_for_existing_database() { let database_identity = test_identity(16); diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index c8c237fa536..ab6e8b7f250 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -1,10 +1,9 @@ use std::time::Duration; -use axum::extract::{Path, State}; +use axum::extract::State; use axum::response::{ErrorResponse, IntoResponse, Response}; use axum::{Extension, Json}; use http::StatusCode; -use serde::Deserialize; use serde_json::{json, Value}; use spacetimedb::auth::identity::ConnectionAuthCtx; use spacetimedb::host::{FunctionArgs, ReducerOutcome}; @@ -15,7 +14,7 @@ use spacetimedb_lib::sats; use super::database::{ client_connected_error_to_response, client_disconnected_error_to_response, find_database_leader, - find_database_module, find_database_or_404, map_reducer_error, sql_direct, SqlQueryParams, + find_database_module, find_database_or_404, map_reducer_error, sql_direct, ResolvedDatabase, SqlQueryParams, }; use crate::auth::SpacetimeAuth; use crate::routes::subscribe::generate_random_connection_id; @@ -38,29 +37,17 @@ const MAX_ERROR_BODY_BYTES: usize = 64 * 1024; type RpcError = (i64, String); -#[derive(Deserialize)] -pub struct McpParams { - name_or_identity: NameOrIdentity, -} - /// handle MCP JSON-RPC request for the database named in the URL -// -// Due to different name resolution and error handling behavior in different branches, -// this route handler does not use [`super::database::resolve_database_name_and_count_response_egress_middleware`]. -// This is unfortunate, as we probably would like to count egress bytes from MCP calls, -// but I (pgoldman 2026-08-07) do not have the wherewithal -// to significantly rewrite this file in order to make it compatible with the middleware, -// and do not know which of its error-handling behaviors are safe to change. pub async fn mcp( State(ctx): State, - Path(McpParams { name_or_identity }): Path, + Extension(ResolvedDatabase(database)): Extension, Extension(auth): Extension, Json(request): Json, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, { - handle_mcp(&ctx, Some(name_or_identity), auth, request).await + handle_mcp(&ctx, Some(database), auth, request).await } pub async fn mcp_root( @@ -76,7 +63,7 @@ where async fn handle_mcp( ctx: &S, - scope: Option, + scope: Option, auth: SpacetimeAuth, request: Value, ) -> axum::response::Result @@ -212,7 +199,7 @@ fn tools_list(host_wide: bool) -> Value { async fn tools_call( ctx: &S, - scope: Option, + scope: Option, auth: SpacetimeAuth, params: Option<&Value>, ) -> Result @@ -260,14 +247,30 @@ where }) } -fn target_database(scope: &Option, arguments: Option<&Value>) -> Result { - if let Some(name_or_identity) = scope { - return Ok(name_or_identity.clone()); +enum Target { + Resolved(Database), + Named(NameOrIdentity), +} + +impl Target { + /// an ErrorResponse here becomes an in-band tool error, an RpcError would be a protocol error + async fn resolve(self, ctx: &(impl ControlStateDelegate + ?Sized)) -> axum::response::Result { + match self { + Target::Resolved(database) => Ok(database), + Target::Named(name_or_identity) => find_database_or_404(ctx, name_or_identity).await, + } + } +} + +fn target_database(scope: &Option, arguments: Option<&Value>) -> Result { + if let Some(database) = scope { + return Ok(Target::Resolved(database.clone())); } let Some(database) = arguments.and_then(|a| a.get("database")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "database argument must be a string".to_owned())); }; serde_json::from_value(Value::String(database.to_owned())) + .map(Target::Named) .map_err(|e| (INVALID_PARAMS, format!("invalid database '{database}': {e}"))) } @@ -319,11 +322,11 @@ where serde_json::to_string(&json!({ "databases": databases })).map_err(log_and_500) } -async fn tool_get_schema(ctx: &S, name_or_identity: NameOrIdentity) -> axum::response::Result +async fn tool_get_schema(ctx: &S, target: Target) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let database = find_database_or_404(ctx, name_or_identity).await?; + let database = target.resolve(ctx).await?; let leader = find_database_leader(ctx, &database).await?; let module = leader.wait_for_module(MODULE_WAIT_TIMEOUT).await.map_err(log_and_500)?; let raw = RawModuleDefV9::from(module.info.module_def.as_ref().clone()); @@ -333,7 +336,7 @@ where async fn tool_sql( ctx: &S, - name_or_identity: NameOrIdentity, + target: Target, auth: SpacetimeAuth, sql: String, confirmed: Option, @@ -343,7 +346,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let database = find_database_or_404(ctx, name_or_identity).await?; + let database = target.resolve(ctx).await?; let rows = sql_direct( ctx.clone(), database, @@ -359,7 +362,7 @@ where async fn tool_call_reducer( ctx: &S, - name_or_identity: NameOrIdentity, + target: Target, auth: SpacetimeAuth, reducer: String, args_json: String, @@ -369,7 +372,7 @@ where { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let database = find_database_or_404(ctx, name_or_identity).await?; + let database = target.resolve(ctx).await?; let module = find_database_module(ctx, &database).await?; let connection_id = generate_random_connection_id(); @@ -541,18 +544,39 @@ mod tests { assert!(get_schema["inputSchema"]["required"].is_null()); } + fn test_database() -> Database { + use spacetimedb::messages::control_db::HostType; + use spacetimedb_lib::Hash; + + Database { + id: 1, + database_identity: Identity::from_byte_array([7; 32]), + owner_identity: Identity::from_byte_array([8; 32]), + host_type: HostType::Wasm, + initial_program: Hash::ZERO, + bootstrap_generation: 0, + } + } + #[test] fn target_database_prefers_the_url_scope_then_the_argument() { - let scoped: NameOrIdentity = serde_json::from_value(json!("mydb")).unwrap(); + let scoped = test_database(); - let target = target_database(&Some(scoped), Some(&json!({ "database": "other" }))).unwrap(); - assert_eq!(target.to_string(), "mydb"); + let scope = Some(scoped.clone()); + let target = target_database(&scope, Some(&json!({ "database": "other" }))).unwrap(); + match target { + Target::Resolved(database) => assert_eq!(database.database_identity, scoped.database_identity), + Target::Named(_) => panic!("expected the database the middleware already resolved"), + } let target = target_database(&None, Some(&json!({ "database": "mydb" }))).unwrap(); - assert_eq!(target.to_string(), "mydb"); + match target { + Target::Named(NameOrIdentity::Name(name)) => assert_eq!(name.as_ref(), "mydb"), + _ => panic!("expected a name"), + } let target = target_database(&None, Some(&json!({ "database": "0".repeat(64) }))).unwrap(); - assert!(matches!(target, NameOrIdentity::Identity(_))); + assert!(matches!(target, Target::Named(NameOrIdentity::Identity(_)))); assert!(target_database(&None, None).is_err()); assert!(target_database(&None, Some(&json!({}))).is_err());