MCP: export_query tool + guard rails for large results

beacon-core/src/extensions.rs

@@ -107,6 +107,12 @@ pub struct McpExtension {
     /// describing what the column means.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub exposed_columns: Option<Vec<ExposedColumn>>,
+    /// Free-form advisory guard rails surfaced to the agent (in the generated
+    /// tool's description and via `describe_table`). Any key/value pairs are
+    /// allowed and beacon does **not** enforce them — they are hints for the model
+    /// (e.g. `{"recommended_row_limit": 10000, "note": "filter by time first"}`).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub guardrails: Option<std::collections::BTreeMap<String, serde_json::Value>>,
 }
 
 /// A column surfaced through the MCP tool — a bare name, or a name plus a

beacon-mcp/src/catalog.rs

@@ -18,7 +18,12 @@ use crate::result::{run_sql_to_json, MAX_ROWS};
 
 /// Build the full tool list: generic tools + per-table tools from extensions.
 pub async fn build_tools(runtime: &Arc<Runtime>) -> anyhow::Result<Vec<Tool>> {
-    let mut tools = vec![list_tables_tool(), describe_table_tool(), run_sql_tool()];
+    let mut tools = vec![
+        list_tables_tool(),
+        describe_table_tool(),
+        run_sql_tool(),
+        export_query_tool(),
+    ];
     for table in runtime.list_tables() {
         let ext = runtime
             .get_table_extensions(table.clone())
@@ -48,6 +53,7 @@ pub async fn dispatch(
                 .ok_or_else(|| anyhow::anyhow!("missing required 'sql' argument"))?;
             run_sql_to_json(runtime, sql.to_string(), identity).await
         }
+        "export_query" => export_query_recipe(&args),
         other => run_table_tool(runtime, other, &args, identity).await,
     }
 }
@@ -92,14 +98,102 @@ fn describe_table_tool() -> Tool {
 fn run_sql_tool() -> Tool {
     read_only(Tool::new(
         "run_sql",
-        "Run a read-only SQL query (SELECT only) against beacon and return JSON rows.",
+        "Run a read-only SQL query (SELECT only) and return JSON rows. This is a bounded \
+         PREVIEW: at most 1000 rows are returned; if the result is larger it is truncated \
+         (the response sets \"truncated\": true). For complete or large results — anything you \
+         intend to analyze in full or hand to a script — use export_query to fetch a \
+         Parquet/Arrow/CSV file instead of run_sql.",
         object_schema(
             json!({ "sql": { "type": "string", "description": "A read-only SELECT statement." } }),
             &["sql"],
         ),
     ))
 }
 
+fn export_query_tool() -> Tool {
+    read_only(Tool::new(
+        "export_query",
+        "Build a recipe to export a large read-only SELECT as a Parquet/Arrow/CSV file for use \
+         in a Python script. Returns the exact /api/query request plus a ready-to-run Python \
+         snippet; it does NOT run the query or return rows. Prefer this over run_sql when the \
+         result is large.",
+        object_schema(
+            json!({
+                "sql": { "type": "string", "description": "A read-only SELECT statement to export." },
+                "format": {
+                    "type": "string",
+                    "enum": ["parquet", "arrow", "csv"],
+                    "description": "Output file format (default parquet)."
+                }
+            }),
+            &["sql"],
+        ),
+    ))
+}
+
+/// Build a "fetch recipe" for exporting a query as a file. MCP tool results are
+/// model-context text, so we never stream the (potentially huge) file through the
+/// model: instead we return the exact `/api/query` request and a Python snippet
+/// the agent can drop into a script, which fetches the Parquet/Arrow/CSV directly.
+fn export_query_recipe(args: &Map<String, Value>) -> anyhow::Result<String> {
+    let sql = args
+        .get("sql")
+        .and_then(Value::as_str)
+        .ok_or_else(|| anyhow::anyhow!("missing required 'sql' argument"))?
+        .trim();
+    // MCP is read-only: only allow SELECT / WITH (CTE) exports.
+    let head = sql.split_whitespace().next().unwrap_or("").to_ascii_uppercase();
+    anyhow::ensure!(
+        matches!(head.as_str(), "SELECT" | "WITH"),
+        "export_query only supports read-only SELECT queries"
+    );
+    let format = args.get("format").and_then(Value::as_str).unwrap_or("parquet");
+    anyhow::ensure!(
+        matches!(format, "parquet" | "arrow" | "csv"),
+        "unsupported format '{format}'; expected one of: parquet, arrow, csv"
+    );
+
+    let body = json!({ "sql": sql, "output": { "format": format } });
+    let body_py = serde_json::to_string(&body)?;
+    let (imports, reader) = match format {
+        "parquet" => ("import io, requests, pandas as pd", "df = pd.read_parquet(io.BytesIO(resp.content))"),
+        "csv" => ("import io, requests, pandas as pd", "df = pd.read_csv(io.BytesIO(resp.content))"),
+        "arrow" => (
+            "import io, requests, pyarrow.ipc as pa_ipc",
+            "df = pa_ipc.open_file(io.BytesIO(resp.content)).read_all().to_pandas()",
+        ),
+        _ => unreachable!(),
+    };
+    let python = [
+        imports.to_string(),
+        "BEACON_URL = \"http://localhost:5001\"  # your beacon host".to_string(),
+        "AUTH = \"Bearer <token>\"  # or \"Basic <base64 user:pass>\"; omit header if anonymous".to_string(),
+        format!(
+            "resp = requests.post(f\"{{BEACON_URL}}/api/query\", headers={{\"Authorization\": AUTH}}, json={body_py})"
+        ),
+        "resp.raise_for_status()".to_string(),
+        reader.to_string(),
+        "print(df.shape)".to_string(),
+    ]
+    .join("\n");
+
+    let recipe = json!({
+        "note": "This does not run the query. POST `request.body` to <BEACON_URL>/api/query; the response body IS the file. Send the same Authorization you use for MCP (Basic/Bearer), or omit it for anonymous access.",
+        "format": format,
+        "request": {
+            "method": "POST",
+            "path": "/api/query",
+            "headers": {
+                "Content-Type": "application/json",
+                "Authorization": "<same credential as MCP; omit if anonymous>"
+            },
+            "body": body
+        },
+        "python": python
+    });
+    Ok(serde_json::to_string_pretty(&recipe)?)
+}
+
 async fn list_tables_json(runtime: &Arc<Runtime>) -> anyhow::Result<String> {
     let mut out = Vec::new();
     for table in runtime.list_tables() {
@@ -218,6 +312,28 @@ fn default_tool_name(table: &str) -> String {
     name
 }
 
+/// Render a table's advisory guard rails to a compact `key: value; …` string for
+/// the tool description. Beacon does not interpret these — they are hints for the
+/// agent. Returns `None` when there are no guard rails.
+fn guardrails_text(mcp: &McpExtension) -> Option<String> {
+    let guardrails = mcp.guardrails.as_ref()?;
+    if guardrails.is_empty() {
+        return None;
+    }
+    let rendered = guardrails
+        .iter()
+        .map(|(key, value)| {
+            let value = value
+                .as_str()
+                .map(str::to_string)
+                .unwrap_or_else(|| value.to_string());
+            format!("{key}: {value}")
+        })
+        .collect::<Vec<_>>()
+        .join("; ");
+    Some(rendered)
+}
+
 async fn table_tool(
     runtime: &Arc<Runtime>,
     table: &str,
@@ -228,10 +344,16 @@ async fn table_tool(
         .tool_name
         .clone()
         .unwrap_or_else(|| default_tool_name(table));
-    let description = mcp
+    let mut description = mcp
         .description
         .clone()
         .unwrap_or_else(|| format!("Query the '{table}' table."));
+    // Surface any advisory guard rails to the model as text (beacon does not
+    // enforce them; they are hints the agent is expected to respect).
+    if let Some(text) = guardrails_text(mcp) {
+        description.push_str("\n\nGuard rails (advisory): ");
+        description.push_str(&text);
+    }
 
     // Merge the table schema (types) with the extension's per-column descriptions,
     // scoped to `exposed_columns` when set, or all columns otherwise, so the model
@@ -445,6 +567,27 @@ mod tests {
         }
     }
 
+    #[test]
+    fn export_query_recipe_builds_fetch_and_guards_writes() {
+        let mut args = Map::new();
+        args.insert("sql".into(), Value::String("SELECT * FROM obs".into()));
+        args.insert("format".into(), Value::String("parquet".into()));
+        let out = export_query_recipe(&args).unwrap();
+        assert!(out.contains("/api/query"), "recipe should reference the query endpoint");
+        assert!(out.contains("read_parquet"), "parquet snippet should use read_parquet");
+        assert!(out.contains("\"format\": \"parquet\""));
+
+        // WITH (CTE) is allowed; default format is parquet.
+        let mut cte = Map::new();
+        cte.insert("sql".into(), Value::String("WITH x AS (SELECT 1) SELECT * FROM x".into()));
+        assert!(export_query_recipe(&cte).unwrap().contains("read_parquet"));
+
+        // Non-SELECT is rejected (MCP is read-only).
+        let mut bad = Map::new();
+        bad.insert("sql".into(), Value::String("DELETE FROM obs".into()));
+        assert!(export_query_recipe(&bad).is_err());
+    }
+
     fn field(name: &str, data_type: &str) -> SchemaFieldView {
         SchemaFieldView {
             name: name.into(),
@@ -480,6 +623,7 @@ mod tests {
                 }),
                 ExposedColumn::Name("lat".into()),
             ]),
+            guardrails: None,
         };
         let cols = resolve_columns(&schema, Some(&ext));
         assert_eq!(cols.len(), 2);
@@ -501,9 +645,25 @@ mod tests {
                     .map(|s| beacon_core::extensions::ExposedColumn::Name(s.to_string()))
                     .collect()
             }),
+            guardrails: None,
         }
     }
 
+    #[test]
+    fn guardrails_render_as_advisory_text() {
+        assert_eq!(guardrails_text(&mcp(None)), None);
+        let mut ext = mcp(None);
+        let mut g = std::collections::BTreeMap::new();
+        g.insert("recommended_row_limit".to_string(), serde_json::json!(10000));
+        g.insert("note".to_string(), serde_json::json!("filter by time first"));
+        ext.guardrails = Some(g);
+        // Rendered as `key: value` pairs (BTreeMap => deterministic key order).
+        assert_eq!(
+            guardrails_text(&ext).as_deref(),
+            Some("note: filter by time first; recommended_row_limit: 10000")
+        );
+    }
+
     #[test]
     fn builds_select_with_preset_between() {
         let p = preset(

beacon-mcp/src/result.rs

@@ -42,8 +42,15 @@ pub async fn run_sql_to_json(
 
     let rows = batches_to_json(&batches)?;
     if truncated {
+        // Guard rail: the result is larger than the preview cap. Return the
+        // preview but steer the model to `export_query` for the complete data
+        // rather than letting it treat the truncated rows as the full result.
         Ok(format!(
-            "{{\"truncated\":true,\"max_rows\":{MAX_ROWS},\"rows\":{rows}}}"
+            "{{\"truncated\":true,\"returned_rows\":{MAX_ROWS},\"max_rows\":{MAX_ROWS},\
+             \"guidance\":\"This is a truncated preview ({MAX_ROWS} rows) because the result is \
+             large. Do NOT treat these rows as complete. To get the full result, call \
+             export_query with the same SQL to fetch it as a Parquet/Arrow/CSV file.\",\
+             \"rows\":{rows}}}"
         ))
     } else {
         Ok(rows)

docs/.vitepress/config.mts

@@ -569,6 +569,20 @@ export default defineConfig({
             },
           ]
         },
+        {
+          text: 'MCP Server',
+          link: '/docs/1.8.0/mcp',
+          collapsed: true,
+          items: [
+            { text: 'Enabling & configuration', link: '/docs/1.8.0/mcp#enabling-configuration' },
+            { text: 'Tools', link: '/docs/1.8.0/mcp#tools' },
+            { text: 'Exposing a table', link: '/docs/1.8.0/mcp#exposing-a-table-to-mcp' },
+            { text: 'Large results', link: '/docs/1.8.0/mcp#large-results' },
+            { text: 'Authenticating an agent', link: '/docs/1.8.0/mcp#authenticating-an-agent' },
+            { text: 'Connecting a client', link: '/docs/1.8.0/mcp#connecting-a-client' },
+            { text: 'How it works', link: '/docs/1.8.0/mcp#how-it-works' },
+          ]
+        },
         {
           text: 'Connect',
           collapsed: true,