fix(catalogs): close resource-consumption gaps in agent catalog access

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<F>(work: F) -> Result<QueryResult, String>
+where
+    F: FnOnce(&Session) -> Result<QueryResult, tower_duckdb::Error> + 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<QueryResult, String> {
     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<QueryResult, tower_duckdb::Error> {
-        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')'";

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<T: serde::Serialize>(data: T) -> Result<CallToolResult, McpError> {
+        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<CallToolResult, McpError> {
         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(),