Add egress tracking directly in mcp_root

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();

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<S>(
 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<S>(
@@ -58,14 +62,23 @@ pub async fn mcp_root<S>(
 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<S>(
     ctx: &S,
     scope: Option<Database>,
     auth: SpacetimeAuth,
     request: Value,
+    // set to the database a tool addressed, so mcp_root can attribute its egress
+    addressed: &mut Option<Identity>,
 ) -> axum::response::Result<Response>
 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<S>(
     scope: Option<Database>,
     auth: SpacetimeAuth,
     params: Option<&Value>,
+    addressed: &mut Option<Identity>,
 ) -> Result<Value, RpcError>
 where
     S: ControlStateDelegate + NodeDelegate + Authorization + Clone + 'static,
@@ -221,22 +235,22 @@ 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)?;
             let Some(reducer) = arguments.and_then(|a| a.get("reducer")).and_then(Value::as_str) else {
                 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<Database> {
-        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<Identity>,
+    ) -> axum::response::Result<Database> {
+        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<S>(ctx: &S, target: Target) -> axum::response::Result<String>
+async fn tool_get_schema<S>(ctx: &S, target: Target, addressed: &mut Option<Identity>) -> axum::response::Result<String>
 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<S>(
     auth: SpacetimeAuth,
     sql: String,
     confirmed: Option<bool>,
+    addressed: &mut Option<Identity>,
 ) -> axum::response::Result<String>
 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<S>(
     auth: SpacetimeAuth,
     reducer: String,
     args_json: String,
+    addressed: &mut Option<Identity>,
 ) -> axum::response::Result<String>
 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();

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::<S>).route_layer(axum::middleware::from_fn_with_state(