From 9f4d8bb2332e5d3bf0a4dfaf50c20738e22e0d69 Mon Sep 17 00:00:00 2001 From: Tien Pham Date: Mon, 24 Aug 2026 18:20:19 +0300 Subject: [PATCH] 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(