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.