From 32cf149b89b213524e1799a2c5d2d0c1764f7989 Mon Sep 17 00:00:00 2001 From: Brad Heller Date: Mon, 27 Jul 2026 15:46:23 +0100 Subject: [PATCH] fix(catalogs): close the resource-consumption gaps in agent catalog access This commit addresses four ways an agent query could spend more than the ceilings intended: The byte ceiling now counts the exact serialized size of each value, including quotes, escapes, and separators. The old estimate counted an array of empty strings as roughly zero bytes, so a nested value like list_transform(range(10000000), x -> '') serialized to 100+ MB while passing a 1 MiB ceiling. The MCP layer also serializes catalog results compactly now, since pretty-printing an array of arrays multiplies its size several-fold past the ceiling the rows were admitted under. Agent-facing error messages are bounded to 4 KiB. A DuckDB error can echo the offending value (a failed CAST reproduces the whole string it was given), which made the error channel an unbounded output path around the result ceilings. Full errors still go to the debug log. DuckDB sessions now run inside a process-wide two-slot budget, and a dropped MCP request interrupts its running query instead of leaving it to burn its ceilings on a thread nobody is waiting on. Without this, concurrent agent calls multiplied the 1 GiB memory / 2 GiB spill ceilings per call. Catalog discovery through MCP runs under the agent ceilings too, with a tables_truncated flag, so a huge catalog cannot flood a model's context through SHOW ALL TABLES. The CLI show path stays unbounded, since a person asked for the listing. --- crates/tower-cmd/src/catalogs.rs | 153 ++++++++++++++--- crates/tower-cmd/src/mcp.rs | 49 ++++-- crates/tower-duckdb/src/lib.rs | 277 ++++++++++++++++++++++++++++--- 3 files changed, 419 insertions(+), 60 deletions(-) diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index 1bc6ea19..0da8428d 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -8,7 +8,7 @@ use tower_api::models::{ catalog_fact, update_catalog_fact_body, vend_catalog_credentials_body, CatalogCredentials, CatalogFact, DescribeCatalogResponse, UpdateCatalogFactBody, }; -use tower_duckdb::{guard, params, run_query, Hardening, Limits, QueryResult, Session}; +use tower_duckdb::{guard, params, CancelHandle, Hardening, Limits, QueryResult, Session}; use tower_telemetry::debug; use crate::{api, beta, output, util::cmd}; @@ -382,7 +382,7 @@ async fn fetch_catalog_tables( let result = if full { list_catalog_columns(config, name, env).await } else { - list_catalog_tables(config, name, env).await + list_catalog_tables(config, name, env, Limits::none()).await }; match result { @@ -397,13 +397,97 @@ async fn fetch_catalog_tables( } } +/// How many DuckDB sessions may run at once across the whole process. +/// +/// Every catalog query and table listing opens its own in-memory session, and a +/// hardened session may spend up to `Hardening::agent()`'s ceilings — 1 GiB of +/// engine memory and 2 GiB of spill — *each*. The MCP server dispatches +/// requests concurrently, so without a shared budget a handful of parallel +/// agent calls multiplies those ceilings until the machine runs out. Two slots +/// keeps one slow scan from blocking every other call outright while capping +/// the worst case at twice the per-session ceilings; everything else queues. +const MAX_CONCURRENT_DUCKDB_SESSIONS: usize = 2; + +static DUCKDB_SESSION_SLOTS: tokio::sync::Semaphore = + tokio::sync::Semaphore::const_new(MAX_CONCURRENT_DUCKDB_SESSIONS); + +/// Cancels the session's query when dropped. Held across the await on the +/// blocking task: if the request future is dropped (an MCP cancellation, a +/// disconnect), the guard's drop interrupts the query instead of leaving it +/// running to its ceilings on a thread nobody is waiting on. Dropping after a +/// normal completion is a no-op. +struct CancelOnDrop(CancelHandle); + +impl Drop for CancelOnDrop { + fn drop(&mut self) { + self.0.cancel(); + } +} + +/// Runs `work` with a fresh DuckDB session on a blocking thread, inside the +/// process-wide session budget and wired for cancellation. Errors come back as +/// strings with no token redaction — the caller holds the token and redacts. +async fn run_bounded_session(work: F) -> Result +where + F: FnOnce(&Session) -> Result + Send + 'static, +{ + let _slot = DUCKDB_SESSION_SLOTS + .acquire() + .await + .expect("session semaphore is never closed"); + + let cancel = CancelHandle::new(); + let _guard = CancelOnDrop(cancel.clone()); + tokio::task::spawn_blocking(move || { + let session = Session::open().map_err(|err| format!("Query failed: {err}"))?; + if cancel.attach(&session) { + // Cancelled while queued for a slot; don't start work nobody awaits. + return Err("The query was cancelled.".to_string()); + } + work(&session).map_err(|err| format!("Query failed: {err}")) + }) + .await + .map_err(|err| format!("Query execution panicked: {err}"))? +} + +/// Ceiling on an error message crossing to an agent. DuckDB errors can echo an +/// offending *value* (a failed CAST reproduces the whole string it was given), +/// which turns the error channel into an unbounded output path around the +/// result ceilings. Generous enough for real diagnostics — parser errors with +/// candidates, binder suggestions — while closing the loophole. +const AGENT_MAX_ERROR_BYTES: usize = 4096; + +/// Bounds an error message to [`AGENT_MAX_ERROR_BYTES`] for the agent-facing +/// paths. The full message (already redacted by the caller) goes to the debug +/// log, so detail is kept where an operator can read it rather than shipped to +/// the model. Truncates on a char boundary. +pub(crate) fn bound_agent_error(message: String) -> String { + if message.len() <= AGENT_MAX_ERROR_BYTES { + return message; + } + debug!("full error before truncation for agent: {message}"); + let mut cut = AGENT_MAX_ERROR_BYTES; + while !message.is_char_boundary(cut) { + cut -= 1; + } + format!( + "{}… [error truncated; {} bytes total]", + &message[..cut], + message.len() + ) +} + /// Attaches a storage catalog read-only and returns its (namespace, table) -/// rows via `SHOW ALL TABLES`. Shared by the CLI `show` and the MCP server. -/// Errors are returned with the OAuth token redacted. +/// rows via `SHOW ALL TABLES`, bounded by `limits`. Shared by the CLI `show` +/// (which passes `Limits::none()` — a person asked for the listing) and the MCP +/// server (which passes `Limits::agent()` and reports truncation, so a huge +/// catalog cannot flood a model's context). Errors are returned with the OAuth +/// token redacted. pub(crate) async fn list_catalog_tables( config: &Config, name: &str, env: &str, + limits: Limits, ) -> Result { let response = api::vend_catalog_credentials(config, name, env, vend_catalog_credentials_body::Mode::Read) @@ -418,17 +502,15 @@ pub(crate) async fn list_catalog_tables( ); let db_name = name.to_string(); - tokio::task::spawn_blocking(move || { - run_query( - &setup, + run_bounded_session(move |session| { + session.run_setup(&setup)?; + session.query( "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", params![db_name], - &Limits::none(), + &limits, ) }) .await - .map_err(|err| err.to_string()) - .and_then(|inner| inner.map_err(|err| err.to_string())) .map_err(|err| redact_token(&err, &token)) } @@ -687,25 +769,17 @@ pub(crate) async fn query_catalog_for_agent( &response.credentials, vend_catalog_credentials_body::Mode::Read, ); - let result = tokio::task::spawn_blocking(move || -> Result { - let session = Session::open()?; + run_bounded_session(move |session| { session.run_setup(&setup)?; session.harden(&Hardening::agent())?; session.query(&sql, [], &Limits::agent()) }) - .await; - - match result { - Ok(Ok(query_result)) => Ok(query_result), - Ok(Err(err)) => Err(format!( - "Query failed: {}", - redact_token(&err.to_string(), &token) - )), - Err(err) => Err(format!( - "Query execution panicked: {}", - redact_token(&err.to_string(), &token) - )), - } + .await + // Redact before bounding, so truncation cannot cut the message ahead of + // the token and leave it intact; bound because a DuckDB error can echo an + // arbitrarily large offending value, which would bypass the result + // ceilings through the error channel. + .map_err(|err| bound_agent_error(redact_token(&err, &token))) } fn read_sql_from_stdin(out: &output::Out) -> String { @@ -1963,6 +2037,35 @@ mod tests { ); } + /// A DuckDB error can echo the offending value — a failed CAST reproduces + /// the whole string it was given — so an unbounded error message is an + /// output channel around the result ceilings. Agent-facing errors are + /// bounded; short ones pass through untouched. + #[test] + fn agent_errors_are_bounded() { + let short = "Conversion Error: Could not convert string 'x' to INT32".to_string(); + assert_eq!(super::bound_agent_error(short.clone()), short); + + let huge = format!( + "Conversion Error: Could not convert string '{}' to INT32", + "x".repeat(2_000_000) + ); + let total = huge.len(); + let bounded = super::bound_agent_error(huge); + assert!( + bounded.len() < super::AGENT_MAX_ERROR_BYTES + 100, + "bounded error is still {} bytes", + bounded.len() + ); + assert!(bounded.starts_with("Conversion Error")); + assert!(bounded.ends_with(&format!("[error truncated; {total} bytes total]"))); + + // Truncation must not split a multi-byte character. + let unicode = "é".repeat(super::AGENT_MAX_ERROR_BYTES); + let bounded = super::bound_agent_error(unicode); + assert!(bounded.contains("[error truncated")); + } + #[test] fn redact_token_scrubs_secret_from_error_text() { let msg = "Parser Error near 'CREATE SECRET tower_cat (TYPE iceberg, TOKEN 'sekret-123')'"; diff --git a/crates/tower-cmd/src/mcp.rs b/crates/tower-cmd/src/mcp.rs index 7ccbce56..a40c48c5 100644 --- a/crates/tower-cmd/src/mcp.rs +++ b/crates/tower-cmd/src/mcp.rs @@ -330,6 +330,21 @@ impl TowerService { Ok(CallToolResult::success(vec![Content::text(text)])) } + /// Like `json_success`, but compact. For the data-carrying catalog results, + /// whose rows are read under a byte ceiling counted in compact JSON: + /// pretty-printing an array of arrays spends several bytes of indentation + /// per value, which would let the serialized response outgrow the ceiling + /// the rows were admitted under. + fn json_success_compact(data: T) -> Result { + let text = serde_json::to_string(&data).map_err(|e| { + McpError::internal_error( + "Serialization failed", + Some(json!({"error": e.to_string()})), + ) + })?; + Ok(CallToolResult::success(vec![Content::text(text)])) + } + fn text_success(message: String) -> Result { Ok(CallToolResult::success(vec![Content::text(message)])) } @@ -732,7 +747,7 @@ impl TowerService { } #[tool( - description = "Show a catalog's details: its property names and, for Tower-managed storage catalogs, the namespaces and tables you can query." + description = "Show a catalog's details: its property names and, for Tower-managed storage catalogs, the namespaces and tables you can query. The table listing is capped; when \"tables_truncated\" is true the catalog has more tables than shown, so query its metadata (e.g. with WHERE filters) to find the rest." )] async fn tower_catalogs_show( &self, @@ -756,14 +771,18 @@ impl TowerService { // Only Tower-managed storage catalogs expose queryable tables; // for anything else `tables` stays null. A listing failure is - // surfaced in `tables_error` without failing the whole call. - let (tables, tables_error) = if crate::catalogs::is_storage_catalog_type(Some( - &catalog.r#type, - )) { - match crate::catalogs::list_catalog_tables( + // surfaced in `tables_error` (bounded, since a DuckDB error can + // be arbitrarily large) without failing the whole call. The + // listing runs under the agent ceilings — this response lands + // in a model's context, so a huge catalog is cut short and + // flagged via `tables_truncated` rather than dumped whole. + let (tables, tables_truncated, tables_error) = + if crate::catalogs::is_storage_catalog_type(Some(&catalog.r#type)) { + match crate::catalogs::list_catalog_tables( &self.config, &request.name, environment, + tower_duckdb::Limits::agent(), ) .await { @@ -780,20 +799,26 @@ impl TowerService { }) .collect(), ), + Value::Bool(result.is_truncated()), + Value::Null, + ), + Err(e) => ( Value::Null, + Value::Null, + Value::String(crate::catalogs::bound_agent_error(e)), ), - Err(e) => (Value::Null, Value::String(e)), } - } else { - (Value::Null, Value::Null) - }; + } else { + (Value::Null, Value::Null, Value::Null) + }; - Self::json_success(json!({ + Self::json_success_compact(json!({ "name": catalog.name, "type": catalog.r#type, "environment": catalog.environment, "properties": properties, "tables": tables, + "tables_truncated": tables_truncated, "tables_error": tables_error, })) } @@ -825,7 +850,7 @@ impl TowerService { ) .await { - Ok(result) => Self::json_success(json!({ + Ok(result) => Self::json_success_compact(json!({ "columns": result.columns, "rows": result.rows, "row_count": result.rows.len(), diff --git a/crates/tower-duckdb/src/lib.rs b/crates/tower-duckdb/src/lib.rs index 57ab1b11..41ccd9c8 100644 --- a/crates/tower-duckdb/src/lib.rs +++ b/crates/tower-duckdb/src/lib.rs @@ -272,25 +272,38 @@ impl Session { truncated = Some(Truncation::Rows); break; } - let mut record = Vec::with_capacity(columns.len()); - for idx in 0..columns.len() { - let value: duckdb::types::Value = row.get(idx)?; - record.push(value_to_json(value)); - } // Measured before the row is kept, so a row that would blow the // budget is discarded rather than returned and merely labelled // truncated. A single `string_agg` can carry an entire table in // one row, so admitting it and flagging it would leave the // ceiling doing nothing at all. A first row that is already over // budget yields an empty, truncated result, which is the honest - // answer. Note this bounds what the caller is handed, not what - // the engine allocated to produce it; that needs - // `Hardening::memory_limit`. - let record_bytes = record.iter().map(json_size).sum::(); - if limits - .max_total_bytes - .is_some_and(|max| total_bytes + record_bytes > max) - { + // answer. The check runs per value, not per row, so conversion + // stops at the first column that exceeds the budget instead of + // materializing the rest of an oversized row. (The offending + // value itself is still built once before it is measured — the + // driver hands values over whole — so the engine's + // `Hardening::memory_limit` is what bounds a query's spend, + // while this bounds what the caller is handed.) + let mut record = Vec::with_capacity(columns.len()); + // Counted as the row serializes: `[`…`]` plus a comma per + // separator, and a comma joining it to the previous row. + let mut record_bytes = 2 + usize::from(!rows.is_empty()); + let mut over_budget = false; + for idx in 0..columns.len() { + let value: duckdb::types::Value = row.get(idx)?; + let value = value_to_json(value); + record_bytes += json_size(&value) + usize::from(idx > 0); + if limits + .max_total_bytes + .is_some_and(|max| total_bytes + record_bytes > max) + { + over_budget = true; + break; + } + record.push(value); + } + if over_budget { truncated = Some(Truncation::Bytes); break; } @@ -367,23 +380,111 @@ impl Drop for Deadline { } } -/// Rough serialized size of a value, used only to bound how much a result may -/// carry back. Strings dominate real results, so they are measured exactly and -/// everything else is approximated. +/// A shareable handle that stops a session's running query from another thread. +/// +/// A session runs its query on a blocking thread, and the async caller that +/// spawned it can go away — an MCP request gets cancelled, its future dropped — +/// with no way to reach the connection. This is that way: create the handle +/// first, hand a clone to the blocking thread to [`attach`] to its session, and +/// [`cancel`] from anywhere. Cancelling before the attach is remembered +/// (`attach` reports it), so work queued behind a concurrency limit can be +/// abandoned before its query ever starts. Cancelling after the session +/// finished or closed is a documented no-op in the driver, so an unconditional +/// cancel-on-drop guard is safe. +/// +/// Like [`Deadline`], this rides on `duckdb_interrupt`, which is honoured at +/// chunk boundaries: it stops a runaway query promptly rather than instantly. +/// +/// [`attach`]: CancelHandle::attach +/// [`cancel`]: CancelHandle::cancel +#[derive(Clone, Default)] +pub struct CancelHandle { + state: std::sync::Arc>, +} + +#[derive(Default)] +struct CancelState { + cancelled: bool, + target: Option>, +} + +impl CancelHandle { + pub fn new() -> Self { + Self::default() + } + + /// Interrupt the attached session's query, if one is running, and remember + /// the cancellation so a later [`attach`](CancelHandle::attach) sees it. + pub fn cancel(&self) { + let target = { + let mut state = self.state.lock().expect("cancel state poisoned"); + state.cancelled = true; + state.target.clone() + }; + if let Some(handle) = target { + handle.interrupt(); + } + } + + /// Bind this handle to `session`'s connection. Returns whether the handle + /// was already cancelled; a caller seeing `true` should drop the session + /// without running its query. + #[must_use] + pub fn attach(&self, session: &Session) -> bool { + let mut state = self.state.lock().expect("cancel state poisoned"); + state.target = Some(session.conn.interrupt_handle()); + state.cancelled + } +} + +/// The exact number of bytes `value` occupies when serialized as compact JSON. +/// +/// This is what the byte ceiling counts, so it must not undercount: an earlier +/// approximation ignored quotes, escapes, and separators, and a query shaped as +/// one deeply nested value (`list_transform(range(10000000), x -> '')`) slipped +/// a 100 MB serialized response under a 1 MiB ceiling. Exactness is asserted +/// against `serde_json::to_string` in the tests. Callers that serialize +/// pretty-printed pay whitespace on top of this, so they should serialize +/// bounded results compactly. fn json_size(value: &serde_json::Value) -> usize { match value { serde_json::Value::Null => 4, - serde_json::Value::Bool(_) => 5, + serde_json::Value::Bool(b) => { + if *b { + 4 + } else { + 5 + } + } serde_json::Value::Number(n) => n.to_string().len(), - serde_json::Value::String(s) => s.len(), - serde_json::Value::Array(items) => items.iter().map(json_size).sum::() + 2, - serde_json::Value::Object(fields) => fields - .iter() - .map(|(key, value)| key.len() + json_size(value)) - .sum::(), + serde_json::Value::String(s) => json_string_size(s), + serde_json::Value::Array(items) => { + 2 + items.len().saturating_sub(1) + items.iter().map(json_size).sum::() + } + serde_json::Value::Object(fields) => { + 2 + fields.len().saturating_sub(1) + + fields + .iter() + .map(|(key, value)| json_string_size(key) + 1 + json_size(value)) + .sum::() + } } } +/// Serialized size of a JSON string: surrounding quotes plus per-byte escape +/// cost, mirroring serde_json's escaping (short escapes for the common control +/// characters, `\u00XX` for the rest, multi-byte UTF-8 passed through). +fn json_string_size(s: &str) -> usize { + 2 + s + .bytes() + .map(|b| match b { + b'"' | b'\\' | 0x08 | 0x0c | b'\n' | b'\r' | b'\t' => 2, + 0x00..=0x1f => 6, + _ => 1, + }) + .sum::() +} + /// Open a session, run `setup`, and execute `query`. Convenience for one-shot /// callers that do not need to hold the session. Callers running untrusted SQL /// should build a [`Session`] and call [`Session::harden`] between setup and @@ -696,6 +797,136 @@ mod tests { ); } + /// The byte counter must be exact for the serialized form, or a value shaped + /// to exploit the gap walks under the ceiling: an array of empty strings used + /// to count ~0 bytes while serializing to 3 bytes per element, letting a + /// 100 MB response through a 1 MiB ceiling. + #[test] + fn json_size_matches_serialized_length_exactly() { + for value in [ + serde_json::json!(null), + serde_json::json!(true), + serde_json::json!(false), + serde_json::json!(0), + serde_json::json!(-12345.678), + serde_json::json!(""), + serde_json::json!("plain"), + serde_json::json!("quote\" backslash\\ newline\n tab\t nul\u{0} unicode\u{1F600}é"), + serde_json::json!([]), + serde_json::json!(["", "", ""]), + serde_json::json!([1, [2, [3, "x\ny"]], {"k": null}]), + serde_json::json!({}), + serde_json::json!({"a": 1, "b\"": ["c", {"d": false}]}), + ] { + let serialized = serde_json::to_string(&value).expect("serialize"); + assert_eq!( + super::json_size(&value), + serialized.len(), + "json_size disagrees with serde_json for {serialized}" + ); + } + } + + /// The nested-value bypass, end to end: one row holding a large list of + /// empty strings must be withheld by the byte ceiling, not returned with + /// `truncated: false`. + #[test] + fn a_nested_value_cannot_walk_under_the_byte_ceiling() { + let result = run_query( + &[], + "SELECT list_transform(range(1000000), x -> '') AS xs", + [], + &Limits { + max_rows: Some(1000), + max_total_bytes: Some(1024 * 1024), + timeout: None, + }, + ) + .expect("query should succeed"); + + assert_eq!(result.truncated, Some(Truncation::Bytes)); + assert!( + result.rows.is_empty(), + "an oversized nested row must be withheld, got {} row(s)", + result.rows.len() + ); + } + + /// Whatever comes back under a byte ceiling must actually serialize inside + /// it — the counted size is the serialized size, so this holds for any + /// result shape, nested or flat. + #[test] + fn returned_rows_serialize_within_the_byte_ceiling() { + let budget = 64 * 1024; + let result = run_query( + &[], + "SELECT list_transform(range(200), x -> 'padding \"quoted\"') AS xs, \ + repeat('y', 500) AS pad FROM range(1000)", + [], + &Limits { + max_rows: None, + max_total_bytes: Some(budget), + timeout: None, + }, + ) + .expect("query should succeed"); + + assert_eq!(result.truncated, Some(Truncation::Bytes)); + let serialized = serde_json::to_string(&result.rows).expect("serialize rows"); + assert!( + serialized.len() <= budget, + "budget was {budget} bytes but rows serialize to {} bytes", + serialized.len() + ); + } + + /// A cancel handle stops a running query from another thread — the path an + /// MCP cancellation takes when the request future is dropped mid-query. + #[test] + fn a_cancel_handle_interrupts_a_running_query() { + let cancel = super::CancelHandle::new(); + let worker_cancel = cancel.clone(); + let started = std::time::Instant::now(); + let worker = std::thread::spawn(move || { + let session = Session::open().expect("open session"); + if worker_cancel.attach(&session) { + panic!("handle cancelled before the query started"); + } + session.query( + "SELECT count(*) FROM range(1000000000000) AS t(i) WHERE i % 7 = 0", + [], + &Limits::none(), + ) + }); + + std::thread::sleep(std::time::Duration::from_millis(300)); + cancel.cancel(); + let result = worker.join().expect("worker should not panic"); + assert!(result.is_err(), "cancelled query should return an error"); + assert!( + started.elapsed() < std::time::Duration::from_secs(60), + "cancel took too long: {:?}", + started.elapsed() + ); + } + + /// Cancelling before the session exists must be remembered: `attach` reports + /// it so a caller queued behind the concurrency limit never starts work for + /// a request that already went away. + #[test] + fn a_cancel_before_attach_is_remembered() { + let cancel = super::CancelHandle::new(); + cancel.cancel(); + let session = Session::open().expect("open session"); + assert!( + cancel.attach(&session), + "attach must report a cancellation that happened first" + ); + // And cancelling after a session is gone is a safe no-op. + drop(session); + cancel.cancel(); + } + /// DuckDB has no statement timeout, so a runaway query is bounded by /// interrupting the connection from the host. `range(1e12)` would run /// effectively forever; it must come back as an error, not hang the test.