Skip to content

MCP: export_query tool + guard rails for large results - #325

Merged
robinskil merged 4 commits into
mainfrom
features/mcp-guardrails
Jul 1, 2026
Merged

robinskil merged 4 commits into
mainfrom
features/mcp-guardrails

Conversation

@robinskil

Copy link
Copy Markdown
Collaborator

Follow-up to the MCP server (#323, now on main). Adds a path for large query results so they don't get streamed through the model's context.

export_query tool

MCP tool results are model-context text, so returning a big Parquet/Arrow blob inline is wrong. export_query instead returns a fetch recipe — it does not run the query:

  • Input: sql (read-only SELECT/WITH) + format (parquet default, arrow IPC, csv).
  • Output: the exact /api/query request (body with output.format) and a ready-to-run Python snippet.
  • A script POSTs that to /api/query, which streams the file in one response — bulk bytes flow script↔beacon, never through the model.
import io, requests, pandas as pd
resp = requests.post(f"{BEACON_URL}/api/query",
    headers={"Authorization": AUTH},
    json={"sql": "SELECT ... FROM obs WHERE ...", "output": {"format": "parquet"}})
resp.raise_for_status()
df = pd.read_parquet(io.BytesIO(resp.content))

No server-side changes — it reuses the existing /api/query file download.

Guard rail on run_sql

run_sql is now explicitly a bounded preview (≤1000 rows). When a result exceeds the cap the response is marked "truncated": true with a guidance field telling the model not to treat the preview as complete and to call export_query for the full result. The tool description states the contract too — so large queries are steered to the file path rather than silently returning partial data.

This is a reactive guard (based on actual rows produced), which is reliable for beacon's streaming/external sources where planner row-count estimates are often absent.

Tests / docs

  • New beacon-mcp unit test for export_query (recipe shape, parquet/CTE, write-rejection); 7 beacon-mcp tests pass, beacon-api builds.
  • docs/mcp.md (tools list + Large results section) and docs/mcp-architecture.md updated.

robinskil added 2 commits July 1, 2026 14:35
run_sql inlines rows (capped) for previews; export_query instead returns a
fetch recipe -- the exact /api/query request plus a ready-to-run Python snippet
-- so a script pulls the Parquet/Arrow/CSV directly from /api/query and the bulk
bytes never enter the model's context. Read-only (SELECT/WITH) guard; parquet
default, arrow (IPC) and csv supported. Tests + docs (user guide + architecture).
run_sql is now explicitly a bounded preview; when a result exceeds the 1000-row
cap the response carries truncated=true plus guidance instructing the model not
to treat the preview as complete and to call export_query for the full result as
a file. Tool description states the contract. Docs updated.
@codecov

codecov Bot commented Jul 1, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 71.07438% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.50%. Comparing base (475e1b5) to head (53ecd03).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
beacon-mcp/src/catalog.rs 74.13% 30 Missing ⚠️
beacon-mcp/src/result.rs 0.00% 5 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main     #325      +/-   ##
==========================================
- Coverage   74.51%   74.50%   -0.01%     
==========================================
  Files         290      290              
  Lines       38448    38566     +118     
==========================================
+ Hits        28649    28735      +86     
- Misses       9799     9831      +32     
Files with missing lines Coverage Δ
beacon-core/src/extensions.rs 88.48% <ø> (ø)
beacon-mcp/src/result.rs 0.00% <0.00%> (ø)
beacon-mcp/src/catalog.rs 55.10% <74.13%> (+5.72%) ⬆️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

robinskil added 2 commits July 1, 2026 15:20
Adds an optional free-form `guardrails` map to the mcp extension: arbitrary
key/value hints that beacon surfaces to the agent (appended to the generated
tool's description and returned by describe_table) but does NOT enforce. Admins
can attach any keys. Enforcement of output size remains the separate built-in
run_sql preview cap. Tests + docs.
Adds a versioned, canonical MCP Server guide at docs/docs/1.8.0/mcp.md (config,
tools incl. export_query, exposing tables via the mcp/preset extensions, column
descriptions, advisory guardrails, large-result export, per-user auth, connecting
Claude, and a 'How it works' section) and wires it into the 1.8.0 sidebar between
REST API and Connect. Removes the redundant top-level docs/mcp.md and
docs/mcp-architecture.md (superseded by the versioned page).
@robinskil robinskil self-assigned this Jul 1, 2026
@robinskil
robinskil merged commit a389bc0 into main Jul 1, 2026
3 checks passed
@robinskil
robinskil deleted the features/mcp-guardrails branch August 12, 2026 15:15
@robinskil robinskil mentioned this pull request Aug 17, 2026
63 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant