From 0d24b961acde63c6c094684b44af555accaa674c Mon Sep 17 00:00:00 2001 From: Tien Pham Date: Fri, 21 Aug 2026 15:22:57 +0300 Subject: [PATCH 1/2] 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()); From fc7b3191d96b12f9ef815b97799821705e3db0c5 Mon Sep 17 00:00:00 2001 From: Tien Pham Date: Mon, 24 Aug 2026 18:20:19 +0300 Subject: [PATCH 2/2] feat: add egress tracking directly in mcp_root --- crates/client-api/src/routes/database.rs | 87 ++++++++++++++++-------- crates/client-api/src/routes/mcp.rs | 57 +++++++++++----- crates/client-api/src/routes/mod.rs | 1 + 3 files changed, 101 insertions(+), 44 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index fffc3ac0b9c..ba7ebcd8008 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -1606,6 +1606,42 @@ where } } +/// Counts a response's headers and body into `spacetime_http_response_size_bytes_total`, +/// attributed to `database_identity` +pub(crate) fn count_response_egress( + database_identity: Identity, + response: axum::response::Response, +) -> axum::response::Response { + let (parts, body) = response.into_parts(); + + // Count the number of bytes used by the headers. + // For guest-defined routes bound to HTTP handlers, these may be arbitrarily large and are worth billing for; + // for built-in routes they will be small and it doesn't really matter one way or another whether we do or don't bill. + // N.b. headers installed by other middleware may or may not be counted here, + // depending on the order in which the middleware applies. + let header_bytes: usize = parts + .headers + .iter() + .map(|(name, value)| name.as_str().len() + value.as_bytes().len()) + .sum(); + + let counter = DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity); + counter.inc_by(header_bytes as u64); + + // `/logs?follow=true` can stream indefinitely. + // Counting frames as they are emitted preserves streaming behavior and avoids buffering the response. + let body = body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + counter.inc_by(data.len() as u64); + } + frame + }); + + axum::response::Response::from_parts(parts, Body::new(body)) +} + /// Resolves an existing database, attaches it as [`ResolvedDatabase`], /// and counts response bytes in the metric `spacetime_http_response_size_bytes_total`. /// @@ -1634,34 +1670,8 @@ where request.extensions_mut().insert(ResolvedDatabase(database.clone())); let response = next.run(request).await; - let (parts, body) = response.into_parts(); - // Count the number of bytes used by the headers. - // For guest-defined routes bound to HTTP handlers, these may be arbitrarily large and are worth billing for; - // for built-in routes they will be small and it doesn't really matter one way or another whether we do or don't bill. - // N.b. headers installed by other middleware may or may not be counted here, - // depending on the order in which the middleware applies. - let header_bytes: usize = parts - .headers - .iter() - .map(|(name, value)| name.as_str().len() + value.as_bytes().len()) - .sum(); - - let counter = DB_METRICS - .http_response_size_bytes - .with_label_values(&database.database_identity); - counter.inc_by(header_bytes as u64); - - // `/logs?follow=true` can stream indefinitely. - // Counting frames as they are emitted preserves streaming behavior and avoids buffering the response. - let body = body.map_frame(move |frame| { - if let Some(data) = frame.data_ref() { - counter.inc_by(data.len() as u64); - } - frame - }); - - Ok(axum::response::Response::from_parts(parts, Body::new(body))) + Ok(count_response_egress(database.database_identity, response)) } #[cfg(test)] @@ -2359,6 +2369,29 @@ mod tests { remove_http_response_size_metric(database_identity); } + #[tokio::test] + async fn count_response_egress_counts_headers_and_body() { + let database_identity = test_identity(21); + remove_http_response_size_metric(database_identity); + + let response = ([(http::header::CONTENT_TYPE, "application/json")], r#"{"ok":true}"#).into_response(); + let counted = count_response_egress(database_identity, response); + + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "application/json".len() as u64 + ); + + let body = counted.into_body().collect().await.unwrap().to_bytes(); + assert_eq!(body, r#"{"ok":true}"#); + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "application/json".len() as u64 + r#"{"ok":true}"#.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(); diff --git a/crates/client-api/src/routes/mcp.rs b/crates/client-api/src/routes/mcp.rs index ab6e8b7f250..2a5c2c156ec 100644 --- a/crates/client-api/src/routes/mcp.rs +++ b/crates/client-api/src/routes/mcp.rs @@ -13,8 +13,9 @@ use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; 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, ResolvedDatabase, SqlQueryParams, + client_connected_error_to_response, client_disconnected_error_to_response, count_response_egress, + find_database_leader, 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; @@ -47,7 +48,10 @@ pub async fn mcp( where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, { - handle_mcp(&ctx, Some(database), auth, request).await + // the middleware counts this route, so the addressed database is discarded + let mut discarded = None; + + handle_mcp(&ctx, Some(database), auth, request, &mut discarded).await } pub async fn mcp_root( @@ -58,7 +62,14 @@ pub async fn mcp_root( where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, { - handle_mcp(&ctx, None, auth, request).await + let mut addressed = None; + let response = handle_mcp(&ctx, None, auth, request, &mut addressed).await?; + + // no path middleware can attribute this route, its database is named in the request body + Ok(match addressed { + Some(database_identity) => count_response_egress(database_identity, response), + None => response, + }) } async fn handle_mcp( @@ -66,6 +77,8 @@ async fn handle_mcp( scope: Option, auth: SpacetimeAuth, request: Value, + // set to the database a tool addressed, so mcp_root can attribute its egress + addressed: &mut Option, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, @@ -85,7 +98,7 @@ where // protocol ping, distinct from the ping tool "ping" => jsonrpc_result(&id, json!({})), "tools/list" => jsonrpc_result(&id, tools_list(host_wide)), - "tools/call" => match tools_call(ctx, scope, auth, request.get("params")).await { + "tools/call" => match tools_call(ctx, scope, auth, request.get("params"), addressed).await { Ok(result) => jsonrpc_result(&id, result), Err((code, message)) => jsonrpc_error(&id, code, message), }, @@ -202,6 +215,7 @@ async fn tools_call( scope: Option, auth: SpacetimeAuth, params: Option<&Value>, + addressed: &mut Option, ) -> Result where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, @@ -221,14 +235,14 @@ where }), // offered only host-wide "list_databases" if scope.is_none() => tool_list_databases(ctx, auth.claims.identity).await, - "get_schema" => tool_get_schema(ctx, target_database(&scope, arguments)?).await, + "get_schema" => tool_get_schema(ctx, target_database(&scope, arguments)?, addressed).await, "sql" => { let target = target_database(&scope, arguments)?; let Some(sql) = arguments.and_then(|a| a.get("sql")).and_then(Value::as_str) else { return Err((INVALID_PARAMS, "sql argument must be a string".to_owned())); }; let confirmed = arguments.and_then(|a| a.get("confirmed")).and_then(Value::as_bool); - tool_sql(ctx, target, auth, sql.to_owned(), confirmed).await + tool_sql(ctx, target, auth, sql.to_owned(), confirmed, addressed).await } "call" => { let target = target_database(&scope, arguments)?; @@ -236,7 +250,7 @@ where return Err((INVALID_PARAMS, "reducer argument must be a string".to_owned())); }; let args_json = reducer_args_json(arguments)?; - tool_call_reducer(ctx, target, auth, reducer.to_owned(), args_json).await + tool_call_reducer(ctx, target, auth, reducer.to_owned(), args_json, addressed).await } other => return Err((INVALID_PARAMS, format!("unknown tool: {other}"))), }; @@ -254,11 +268,18 @@ enum Target { 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, - } + async fn resolve( + self, + ctx: &(impl ControlStateDelegate + ?Sized), + addressed: &mut Option, + ) -> axum::response::Result { + let database = match self { + Target::Resolved(database) => database, + Target::Named(name_or_identity) => find_database_or_404(ctx, name_or_identity).await?, + }; + *addressed = Some(database.database_identity); + + Ok(database) } } @@ -322,11 +343,11 @@ where serde_json::to_string(&json!({ "databases": databases })).map_err(log_and_500) } -async fn tool_get_schema(ctx: &S, target: Target) -> axum::response::Result +async fn tool_get_schema(ctx: &S, target: Target, addressed: &mut Option) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate, { - let database = target.resolve(ctx).await?; + let database = target.resolve(ctx, addressed).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()); @@ -340,13 +361,14 @@ async fn tool_sql( auth: SpacetimeAuth, sql: String, confirmed: Option, + addressed: &mut Option, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let database = target.resolve(ctx).await?; + let database = target.resolve(ctx, addressed).await?; let rows = sql_direct( ctx.clone(), database, @@ -366,13 +388,14 @@ async fn tool_call_reducer( auth: SpacetimeAuth, reducer: String, args_json: String, + addressed: &mut Option, ) -> axum::response::Result where S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static, { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); - let database = target.resolve(ctx).await?; + let database = target.resolve(ctx, addressed).await?; let module = find_database_module(ctx, &database).await?; let connection_id = generate_random_connection_id(); diff --git a/crates/client-api/src/routes/mod.rs b/crates/client-api/src/routes/mod.rs index 14bd3fd810a..19b92e19402 100644 --- a/crates/client-api/src/routes/mod.rs +++ b/crates/client-api/src/routes/mod.rs @@ -37,6 +37,7 @@ where .nest("/energy", energy::router()) .nest("/prometheus", prometheus::router()) .nest("/metrics", metrics::router()) + // 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(