feat(mcp): make dbx_execute_query row limit configurable

crates/dbx-core/src/ai/agent_tools.rs

@@ -27,9 +27,23 @@ const SAMPLE_DATA_LIMIT: usize = 20;
 /// Maximum number of rows returned by browse_collection tool.
 const BROWSE_COLLECTION_LIMIT: usize = 20;
 
-/// Absolute maximum rows any query tool may request.
+/// Absolute maximum rows requested by the sampling tools (get_sample_data,
+/// browse_collection) and by execute_query on MongoDB shell commands.
 const MAX_ALLOWED_ROWS: usize = 100;
 
+/// Absolute maximum rows `execute_query` may request on SQL connections.
+/// MCP publishes this as the `max_rows` tool parameter; kept below the driver
+/// fetch ceiling (`dbx_drivers::execution::MAX_ROWS = 10000`).
+pub const MAX_EXECUTE_QUERY_ROWS: usize = 1_000;
+
+/// The published "up to `MAX_EXECUTE_QUERY_ROWS` rows" contract only holds while this
+/// ceiling stays below what the driver will actually fetch. If the agent ceiling ever met
+/// or exceeded the driver ceiling, a large `max_rows` request would come back truncated by
+/// the driver instead of by our own cap, and the tool description would be a lie.
+///
+/// Asserted at compile time so it cannot be silently dropped.
+const _: () = assert!(MAX_EXECUTE_QUERY_ROWS < dbx_drivers::execution::MAX_ROWS);
+
 /// Default string-cell character budget for AI and local MCP query results.
 const DEFAULT_QUERY_CELL_CHAR_LIMIT: usize = 200;
 
@@ -131,6 +145,15 @@ fn bounded_query_cell_option(value: Option<u64>, default: usize, maximum: usize)
     value.map(|value| value.min(maximum as u64) as usize).unwrap_or(default)
 }
 
+/// Resolve the row cap for execute_query. Splitting by db_type is deliberate:
+/// SQL connections may go up to MAX_EXECUTE_QUERY_ROWS, while the MongoDB shell
+/// path keeps MAX_ALLOWED_ROWS, so widening the SQL ceiling does not also
+/// widen what a Mongo command can return.
+fn execute_query_row_limit(requested: Option<usize>, db_type: &DatabaseType) -> usize {
+    let ceiling = if *db_type == DatabaseType::MongoDb { MAX_ALLOWED_ROWS } else { MAX_EXECUTE_QUERY_ROWS };
+    requested.map(|rows| rows.min(ceiling)).unwrap_or(EXECUTE_QUERY_LIMIT.min(ceiling))
+}
+
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct AgentSqlPermissions {
     pub allow_writes: bool,
@@ -441,7 +464,7 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition {
     } else if sql_permissions.allow_writes {
         "Execute SQL after the user explicitly confirmed this operation. Read queries and non-DDL writes are allowed for this run."
     } else {
-        "Execute a read-only SQL query and return results (max 50 rows). Cross-database reads may use fully qualified names such as database.table (or database.schema.table for SQL Server) without switching the current database. This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run."
+        "Execute a read-only SQL query and return results (default 50 rows, up to 1000 with the limit argument). Cross-database reads may use fully qualified names such as database.table (or database.schema.table for SQL Server) without switching the current database. This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run."
     };
     ToolDefinition {
         name: "execute_query",
@@ -455,7 +478,9 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition {
                 },
                 "limit": {
                     "type": "number",
-                    "description": "Max rows to return (default 50, max 100)"
+                    "minimum": 1,
+                    "maximum": MAX_EXECUTE_QUERY_ROWS,
+                    "description": format!("Max rows to return (default {EXECUTE_QUERY_LIMIT}, max {MAX_EXECUTE_QUERY_ROWS})")
                 },
                 "cell_char_offset": {
                     "type": "integer",
@@ -840,12 +865,8 @@ async fn execute_execute_query(
         return Err("SQL query cannot be empty".to_string());
     }
 
-    let limit = tool_call
-        .arguments
-        .get("limit")
-        .and_then(|v| v.as_u64())
-        .map(|l| (l as usize).min(MAX_ALLOWED_ROWS))
-        .unwrap_or(EXECUTE_QUERY_LIMIT);
+    let limit =
+        execute_query_row_limit(tool_call.arguments.get("limit").and_then(|v| v.as_u64()).map(|l| l as usize), db_type);
     let cell_window = QueryCellWindow::from_arguments(&tool_call.arguments);
 
     if *db_type == DatabaseType::MongoDb {
@@ -988,9 +1009,18 @@ pub fn format_query_result_as_text(
         lines.push(format!("| {} |", cells.join(" | ")));
     }
 
-    // Truncation notice
+    // Truncation notice. `truncated` is not only a row-cap signal: a driver-side
+    // size budget (e.g. the SQLite worker's blob/response limits) raises it
+    // before `limit` rows come back, so only claim the row cap when the returned
+    // rows actually reach it. The remedy differs — raising the limit helps in one
+    // case and does nothing in the other — so the notice never states one.
     if result.truncated || result.rows.len() >= limit {
-        lines.push(format!("... (showing {} rows, result may be truncated)", result.rows.len()));
+        let detail = if result.rows.len() >= limit {
+            format!("the {limit}-row cap was reached — the result may be truncated")
+        } else {
+            format!("the result was truncated by a driver size limit, not the {limit}-row cap")
+        };
+        lines.push(format!("... (showing {} rows; {detail})", result.rows.len()));
     }
 
     // Stats line
@@ -1377,6 +1407,41 @@ async fn resolve_chroma_collection_uuid(
 mod tests {
     use super::*;
 
+    #[test]
+    fn execute_query_row_limit_splits_sql_and_mongo_ceilings() {
+        // No explicit limit keeps the historical default.
+        assert_eq!(execute_query_row_limit(None, &DatabaseType::Postgres), EXECUTE_QUERY_LIMIT);
+        // SQL connections may go up to the published ceiling.
+        assert_eq!(execute_query_row_limit(Some(500), &DatabaseType::Postgres), 500);
+        assert_eq!(execute_query_row_limit(Some(100_000), &DatabaseType::Postgres), MAX_EXECUTE_QUERY_ROWS);
+        // 0 keeps today's pass-through behaviour (the MCP layer clamps to 1).
+        assert_eq!(execute_query_row_limit(Some(0), &DatabaseType::Postgres), 0);
+        // MongoDB keeps its own, narrower ceiling.
+        assert_eq!(execute_query_row_limit(Some(500), &DatabaseType::MongoDb), MAX_ALLOWED_ROWS);
+        assert_eq!(execute_query_row_limit(None, &DatabaseType::MongoDb), EXECUTE_QUERY_LIMIT);
+    }
+
+    #[test]
+    fn execute_query_tool_schema_publishes_the_sql_row_ceiling_only() {
+        let tool = |db_type| {
+            all_tools(db_type, AgentSqlPermissions::default())
+                .into_iter()
+                .find(|tool| tool.name == "execute_query")
+                .expect("execute_query must be registered for this database type")
+        };
+
+        // The published SQL schema must track the single ceiling constant, or the
+        // model would be told a bound the backend does not enforce.
+        assert_eq!(
+            tool(DatabaseType::Postgres).parameters["properties"]["limit"]["maximum"].as_u64(),
+            Some(MAX_EXECUTE_QUERY_ROWS as u64)
+        );
+        // The MongoDB variant has its own tool definition and must stay at 100.
+        let mongo_limit = &tool(DatabaseType::MongoDb).parameters["properties"]["limit"];
+        assert_eq!(mongo_limit["maximum"].as_u64(), None);
+        assert!(mongo_limit["description"].as_str().unwrap_or_default().contains("max 100"));
+    }
+
     #[test]
     fn agent_query_timeout_resolves_by_precedence() {
         // Explicit per-call value wins.
@@ -1848,6 +1913,32 @@ for line in sys.stdin:
         );
     }
 
+    #[test]
+    fn query_result_formatter_only_claims_the_row_cap_when_it_was_reached() {
+        let row = vec![vec![serde_json::json!(1)]];
+
+        // Returned rows reach the cap: the notice names the effective cap so the
+        // model can connect it to the tool's limit argument.
+        let at_cap = query_result(vec!["id"], row.clone(), 0);
+        let output = format_query_result_as_text(&at_cap, 1, QueryCellWindow::default()).unwrap();
+        assert!(
+            output.contains("... (showing 1 rows; the 1-row cap was reached — the result may be truncated)"),
+            "unexpected notice: {output}"
+        );
+
+        // `truncated` without reaching the cap comes from a driver size budget
+        // (SQLite worker blobs / response size). It must not claim the row cap,
+        // because raising the limit would not return the missing rows.
+        let mut size_limited = query_result(vec!["id"], row, 0);
+        size_limited.truncated = true;
+        let output = format_query_result_as_text(&size_limited, 100, QueryCellWindow::default()).unwrap();
+        assert!(
+            output.contains("the result was truncated by a driver size limit, not the 100-row cap"),
+            "unexpected notice: {output}"
+        );
+        assert!(!output.contains("cap was reached"), "a size truncation must not claim the row cap: {output}");
+    }
+
     #[test]
     fn query_result_formatter_appends_server_messages() {
         let mut dml = query_result(vec![], vec![], 2);

crates/dbx-mcp/src/backend.rs

@@ -1226,8 +1226,17 @@ impl DbxBackend for WebBackend {
                 }
             }
 
-            let max_rows = arguments.get("limit").and_then(Value::as_u64).unwrap_or(100) as usize;
-            let mut body = json!({ "connectionId": connection.id, "database": database, "sql": sql });
+            // Clamp here too: the Web backend does not go through the in-process
+            // `execute_tool` clamp, so an unclamped value would bypass the
+            // published max_rows ceiling. `maxRows` also bounds how many rows the
+            // route fetches, not just how many are rendered.
+            let max_rows = arguments
+                .get("limit")
+                .and_then(Value::as_u64)
+                .unwrap_or(100)
+                .clamp(1, agent_tools::MAX_EXECUTE_QUERY_ROWS as u64) as usize;
+            let mut body =
+                json!({ "connectionId": connection.id, "database": database, "sql": sql, "maxRows": max_rows });
             // Stateful MCP sessions pin every query to the same backend pool.
             if let Some(client_session_id) = arguments.get("client_session_id").and_then(Value::as_str) {
                 body["clientSessionId"] = json!(client_session_id);
@@ -2700,14 +2709,16 @@ mod tests {
         assert_eq!(request_line, "POST /api/query/execute HTTP/1.1");
         let request: Value = serde_json::from_str(&body).unwrap();
         assert_eq!(request["timeoutSecs"], 60);
+        assert_eq!(request["maxRows"], 10);
 
-        // Policy argument 300 overrides the connection.
+        // Policy argument 300 overrides the connection; maxRows is clamped to the
+        // published ceiling instead of being forwarded as-is.
         let result = backend
             .execute_agent_tool(
                 &connection,
                 "postgres",
                 "execute_query",
-                json!({ "sql": "SELECT 1", "limit": 10, "timeout_secs": 300 }),
+                json!({ "sql": "SELECT 1", "limit": 100000, "timeout_secs": 300 }),
                 AgentSqlPermissions { allow_writes: false, allow_dangerous: false, confirmed_write_sql: None },
             )
             .await;
@@ -2717,6 +2728,7 @@ mod tests {
         let (_request_line, second_body) = request_receiver.recv().unwrap();
         let second_request: Value = serde_json::from_str(&second_body).unwrap();
         assert_eq!(second_request["timeoutSecs"], 300);
+        assert_eq!(second_request["maxRows"], agent_tools::MAX_EXECUTE_QUERY_ROWS);
     }
 
     #[cfg(feature = "mq-admin")]

crates/dbx-mcp/src/server.rs

@@ -13,7 +13,9 @@ use crate::backend::{format_query_result, new_connection_config, parse_database_
 use crate::mongo::{self, MongoCommand, MongoSafetyError};
 use crate::session::{McpSession, McpSessionStore};
 use dbx_core::{
-    agent_tools::{format_query_result_as_text, normalize_sql_for_confirmation, QueryCellWindow},
+    agent_tools::{
+        format_query_result_as_text, normalize_sql_for_confirmation, QueryCellWindow, MAX_EXECUTE_QUERY_ROWS,
+    },
     database_manifest,
     db::redis_driver::{classify_command, parse_command_argv, RedisCommandResult, RedisCommandSafety},
     history::HistoryEntry,
@@ -135,6 +137,11 @@ pub struct ExecuteQueryRequest {
     )]
     #[schemars(extend("type" = "integer"))]
     pub cell_char_limit: Option<u64>,
+    #[schemars(
+        description = "Maximum rows returned for this call (default 100, max 1000; larger values are clamped). SQL connections only - MongoDB shell commands always return at most 100 rows, and multi-statement scripts routed to the batch executor return at most 100 rows per statement."
+    )]
+    #[schemars(extend("type" = "integer"))]
+    pub max_rows: Option<u64>,
 }
 
 #[derive(Debug, Deserialize, schemars::JsonSchema)]
@@ -639,7 +646,7 @@ impl DbxMcpServer {
 
     #[tool(
         name = "dbx_execute_query",
-        description = "Execute a SQL query on a database connection (max 100 rows returned). For backwards compatibility, multi-statement scripts and stored-procedure scripts are routed through the dialect-aware batch executor."
+        description = "Execute a SQL query on a database connection (default 100 rows; pass max_rows to return up to 1000 rows; larger values are clamped; MongoDB shell commands and multi-statement scripts routed through the batch executor stay at 100 rows). For backwards compatibility, multi-statement scripts and stored-procedure scripts are routed through the dialect-aware batch executor."
     )]
     async fn execute_query(&self, Parameters(request): Parameters<ExecuteQueryRequest>) -> CallToolResult {
         if let Err(error) = self.ensure_tool_allowed("dbx_execute_query").await {
@@ -721,6 +728,9 @@ impl DbxMcpServer {
             }
         }
         if connection.db_type == DatabaseType::MongoDb {
+            // max_rows is deliberately ignored here: MongoDB shell commands
+            // always return at most 100 rows, and this branch returns before the
+            // arguments below are built.
             let command = match validate_mongo_command_with_groups(
                 connection,
                 &resolved.policy,
@@ -763,8 +773,9 @@ impl DbxMcpServer {
             Ok(permissions) => permissions,
             Err(error) => return error,
         };
+        let max_rows = request.max_rows.map(|value| value.clamp(1, MAX_EXECUTE_QUERY_ROWS as u64)).unwrap_or(100);
         let history_sql = request.sql.clone();
-        let mut arguments = json!({ "sql": request.sql, "limit": 100 });
+        let mut arguments = json!({ "sql": request.sql, "limit": max_rows });
         if let Some(schema) = self.scope.schema.as_deref() {
             arguments["schema"] = json!(schema);
         }
@@ -1929,9 +1940,9 @@ fn sql_requires_batch_execution(sql: &str, database_type: DatabaseType) -> bool
     dbx_core::sql::sql_execution_plan_for_database(sql, database_type).statements.len() > 1
 }
 
-/// Maximum rows returned per statement in a `dbx_execute_batch` call, matching
-/// the `dbx_execute_query` "at most 100 rows" contract so batch responses stay
-/// bounded.
+/// Maximum rows returned per statement in a `dbx_execute_batch` call. The batch
+/// executor is deliberately tighter than dbx_execute_query's default: a script
+/// has many statements, so its row budget multiplies.
 const BATCH_MAX_ROWS: usize = 100;
 
 fn mcp_sql_activity_kind(sql: &str, database_type: DatabaseType) -> &'static str {
@@ -3043,7 +3054,7 @@ mod tests {
             ("dbx_describe_table", &["database", "schema"]),
             ("dbx_list_routines", &["database", "schema", "routine_type"]),
             ("dbx_get_routine_source", &["database", "schema", "signature"]),
-            ("dbx_execute_query", &["database", "session_id", "cell_char_offset", "cell_char_limit"]),
+            ("dbx_execute_query", &["database", "session_id", "cell_char_offset", "cell_char_limit", "max_rows"]),
             ("dbx_execute_batch", &["database", "session_id", "continue_on_error", "use_transaction"]),
             ("dbx_open_session", &["database"]),
             ("dbx_open_table", &["database", "schema"]),
@@ -3100,11 +3111,15 @@ mod tests {
 
         assert!(omitted.selector.connection_id.is_none());
         assert!(omitted.selector.connection_name.is_none());
+        assert!(omitted.max_rows.is_none());
         assert!(explicit_nulls.selector.connection_id.is_none());
         assert!(explicit_nulls.selector.connection_name.is_none());
         assert_eq!(by_name.selector.connection_name.as_deref(), Some("test_conn"));
         assert_eq!(by_id.selector.connection_id.as_deref(), Some("123e4567-e89b-12d3-a456-426614174000"));
 
+        let with_max_rows: ExecuteQueryRequest = serde_json::from_str(r#"{"sql":"SELECT 1","max_rows":500}"#).unwrap();
+        assert_eq!(with_max_rows.max_rows, Some(500));
+
         let nested = serde_json::from_str::<ExecuteQueryRequest>(
             r#"{"connection_name":{"tool":"dbx_dbx_execute_query","error":"Invalid input"},"sql":"SELECT 1"}"#,
         )
@@ -3338,6 +3353,7 @@ mod tests {
                 session_id: None,
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert_eq!(result_text(&show_databases), "- aa\n- aaa\n- abc");
@@ -3735,6 +3751,7 @@ mod tests {
                 session_id: Some(session_id.clone()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert_eq!(result_text(&result), "ok");
@@ -3754,6 +3771,7 @@ mod tests {
                 session_id: Some(session_id.clone()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert!(result_text(&mismatch).contains("SESSION_DATABASE_MISMATCH"));
@@ -3772,6 +3790,7 @@ mod tests {
                 session_id: Some(session_id.clone()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert!(result_text(&missing).contains("SESSION_NOT_FOUND"));
@@ -3794,6 +3813,7 @@ mod tests {
                 session_id: None,
                 cell_char_offset: Some(200),
                 cell_char_limit: Some(800),
+                max_rows: None,
             }))
             .await;
 
@@ -3804,6 +3824,38 @@ mod tests {
         assert_eq!(arguments["cell_char_limit"], 800);
     }
 
+    #[tokio::test]
+    async fn execute_query_forwards_and_clamps_max_rows() {
+        let elasticsearch = connection("esm", "esm", "elasticsearch", "");
+        let backend = Arc::new(FakeBackend { connections: vec![elasticsearch], ..Default::default() });
+        let server = DbxMcpServer::with_runtime_options(backend.clone(), McpScope::default(), false);
+
+        for max_rows in [None, Some(500), Some(100_000), Some(0)] {
+            let result = server
+                .execute_query(Parameters(ExecuteQueryRequest {
+                    selector: selector("esm"),
+                    database: None,
+                    sql: "GET /logs/_search".to_string(),
+                    session_id: None,
+                    cell_char_offset: None,
+                    cell_char_limit: None,
+                    max_rows,
+                }))
+                .await;
+            assert_eq!(result_text(&result), "ok");
+        }
+
+        let recorded = backend.recorded_arguments.lock().unwrap();
+        let limits = recorded
+            .iter()
+            .filter(|(name, _)| name == "execute_query")
+            .map(|(_, arguments)| arguments["limit"].as_u64().expect("limit should be published"))
+            .collect::<Vec<_>>();
+        // Omitted keeps the historical 100-row default; explicit values are clamped
+        // into 1..=MAX_EXECUTE_QUERY_ROWS rather than rejected.
+        assert_eq!(limits, vec![100, 500, MAX_EXECUTE_QUERY_ROWS as u64, 1]);
+    }
+
     #[tokio::test]
     async fn expired_sessions_close_backend_pools_without_accumulating() {
         let starrocks = connection("sr", "sr", "starrocks", "default_catalog");
@@ -3823,6 +3875,7 @@ mod tests {
                     session_id: Some(session_id.clone()),
                     cell_char_offset: None,
                     cell_char_limit: None,
+                    max_rows: None,
                 }))
                 .await;
             assert_eq!(result_text(&result), "ok");
@@ -3860,6 +3913,7 @@ mod tests {
                 session_id: Some(session_id.clone()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert_eq!(result_text(&query), "ok");
@@ -3876,6 +3930,7 @@ mod tests {
                 session_id: Some(session_id.clone()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert_eq!(result_text(&retry_query), "ok");
@@ -3907,6 +3962,7 @@ mod tests {
                 session_id: Some("mcp-session-nope".to_string()),
                 cell_char_offset: None,
                 cell_char_limit: None,
+                max_rows: None,
             }))
             .await;
         assert!(result_text(&missing).contains("SESSION_NOT_FOUND"));
@@ -4056,6 +4112,7 @@ mod tests {
             database: None,
             sql: sql.to_string(),
             session_id: None,
+            max_rows: None,
             cell_char_offset: None,
             cell_char_limit: None,
         };

crates/dbx-mcp/tests/protocol.rs

@@ -249,6 +249,30 @@ async fn execute_query_injects_and_omits_timeout_secs_from_policy() {
     );
 }
 
+/// The client-facing field is `max_rows`. Only this rmcp-level test proves the
+/// published name survives tool-call serialization and that clamping happens at
+/// the native boundary; the in-process tests assert the internal `limit`
+/// argument directly and cannot catch a name mismatch.
+#[tokio::test]
+async fn execute_query_forwards_client_max_rows() {
+    let cases = [
+        (json!({ "connection_id": "scoped", "sql": "SELECT 1" }), 100_u64),
+        (json!({ "connection_id": "scoped", "sql": "SELECT 1", "max_rows": 500 }), 500),
+        (json!({ "connection_id": "scoped", "sql": "SELECT 1", "max_rows": 100000 }), 1000),
+        (json!({ "connection_id": "scoped", "sql": "SELECT 1", "max_rows": 0 }), 1),
+    ];
+    for (request, expected) in cases {
+        let backend = Arc::new(CapturingBackend {
+            policy: McpGlobalPolicy::default(),
+            connections: vec![test_connection("scoped", "shared-db")],
+            calls: Mutex::new(Vec::new()),
+        });
+        let captured = captured_query_arguments(backend, request.clone()).await;
+        assert_eq!(captured.len(), 1, "expected one captured execute_query call");
+        assert_eq!(captured[0]["limit"], json!(expected), "max_rows in {request} must map to limit");
+    }
+}
+
 fn test_connection(id: &str, name: &str) -> ConnectionConfig {
     serde_json::from_value(json!({
         "id": id,