Skip to main content

stackql_mcp/
lib.rs

1//! Embedded StackQL MCP server for Rust agentic apps.
2//!
3//! StackQL exposes cloud providers (AWS, GitHub, Google, Azure, ...) as SQL
4//! tables, served over the Model Context Protocol. This crate acquires the
5//! `stackql` binary, launches it as an MCP server over stdio, and hands you a
6//! connected [`rmcp`] client.
7//!
8//! Two acquisition modes behind one API:
9//!
10//! - sidecar (default feature): download the platform's .mcpb bundle at first
11//!   run, verify its sha256 against pins baked into the crate, and cache it
12//!   under `~/.stackql/mcp-server-bin/` (shared with the npm and PyPI
13//!   wrappers)
14//! - vendored (`vendored` feature): embed the .mcpb with `include_bytes!` and
15//!   extract on first run - no network at runtime, single shippable binary
16//!
17//! ```no_run
18//! use stackql_mcp::{Mode, StackqlMcp};
19//!
20//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
21//! let server = StackqlMcp::builder()
22//!     .mode(Mode::ReadOnly)
23//!     .auth(serde_json::json!({"github": {"type": "null_auth"}}))
24//!     .start()
25//!     .await?;
26//! let tools = server.list_all_tools().await?;
27//! println!("{} tools available", tools.len());
28//! server.shutdown().await?;
29//! # Ok(())
30//! # }
31//! ```
32
33mod acquire;
34mod bundle;
35mod cache;
36mod download;
37mod error;
38mod launch;
39mod pins;
40mod platform;
41
42use std::ops::Deref;
43use std::path::PathBuf;
44use std::process::Stdio;
45
46use rmcp::service::RunningService;
47use rmcp::{RoleClient, ServiceExt};
48
49pub use cache::{ENV_BIN, ENV_BUNDLE};
50pub use error::{Error, Result};
51pub use pins::{Pin, PINS, STACKQL_VERSION};
52pub use platform::Platform;
53
54/// Safety contract for query / mutation / lifecycle tools, enforced
55/// server-side. Maps to `server.mode` in the server's `--mcp.config`.
56#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
57pub enum Mode {
58    /// SELECT and metadata tools only. The default: escalation is an
59    /// explicit caller opt-in.
60    #[default]
61    ReadOnly,
62    /// Reads plus non-destructive mutations (the server's own default).
63    Safe,
64    /// Safe plus deletes.
65    DeleteSafe,
66    /// All operations, including lifecycle provisioning.
67    FullAccess,
68}
69
70impl Mode {
71    /// The wire value for `server.mode`.
72    pub fn as_str(self) -> &'static str {
73        match self {
74            Mode::ReadOnly => "read_only",
75            Mode::Safe => "safe",
76            Mode::DeleteSafe => "delete_safe",
77            Mode::FullAccess => "full_access",
78        }
79    }
80}
81
82/// Download the pinned .mcpb bundle for the host platform into the shared
83/// cache (verified against the baked sha256 pin) and return its path. Skips
84/// the download when a verified copy is already present.
85///
86/// This is the producer side of vendored builds: fetch the bundle once on the
87/// build machine, then embed it with [`include_bundle!`].
88#[cfg(feature = "sidecar")]
89pub fn fetch_bundle() -> Result<PathBuf> {
90    let platform = Platform::detect()?;
91    let pin = pins::pin_for(platform)?;
92    let dest = cache::bin_cache_root()?
93        .join(pins::STACKQL_VERSION)
94        .join(pin.bundle_name);
95    if dest.is_file() && download::sha256_file(&dest)? == pin.sha256 {
96        return Ok(dest);
97    }
98    download::download_verified(&pins::bundle_url(pin), pin.sha256, &dest)?;
99    Ok(dest)
100}
101
102/// Embed the .mcpb bundle named by the compile-time env var
103/// `STACKQL_MCP_BUNDLE_FILE`, for use with `Builder::bundle_bytes` (vendored
104/// feature):
105///
106/// ```ignore
107/// let server = StackqlMcp::builder()
108///     .bundle_bytes(stackql_mcp::include_bundle!())
109///     .start()
110///     .await?;
111/// ```
112///
113/// Build with `STACKQL_MCP_BUNDLE_FILE=/abs/path/to/bundle.mcpb cargo build`.
114/// Pair with [`fetch_bundle`] to produce the bundle.
115#[macro_export]
116macro_rules! include_bundle {
117    () => {
118        include_bytes!(env!(
119            "STACKQL_MCP_BUNDLE_FILE",
120            "set STACKQL_MCP_BUNDLE_FILE to the absolute path of the platform .mcpb bundle \
121             (see stackql_mcp::fetch_bundle)"
122        ))
123    };
124}
125
126/// Entry point. See the crate docs for the full example.
127pub struct StackqlMcp;
128
129impl StackqlMcp {
130    pub fn builder() -> Builder {
131        Builder::default()
132    }
133}
134
135/// Configures and starts the embedded server.
136#[derive(Default)]
137pub struct Builder {
138    mode: Mode,
139    auth: Option<serde_json::Value>,
140    approot: Option<PathBuf>,
141    acquisition: acquire::Acquisition,
142}
143
144impl Builder {
145    /// Safety mode for the server. Defaults to [`Mode::ReadOnly`].
146    pub fn mode(mut self, mode: Mode) -> Self {
147        self.mode = mode;
148        self
149    }
150
151    /// Provider auth document, passed to the server as `--auth=<json>`.
152    /// Example: `json!({"github": {"type": "null_auth"}})`.
153    pub fn auth(mut self, auth: serde_json::Value) -> Self {
154        self.auth = Some(auth);
155        self
156    }
157
158    /// Override the server's application root. Defaults to `<home>/.stackql`.
159    pub fn approot(mut self, approot: impl Into<PathBuf>) -> Self {
160        self.approot = Some(approot.into());
161        self
162    }
163
164    /// Run an existing stackql binary instead of acquiring one. The
165    /// `STACKQL_MCP_BIN` env var takes precedence over this.
166    pub fn binary(mut self, path: impl Into<PathBuf>) -> Self {
167        self.acquisition.binary = Some(path.into());
168        self
169    }
170
171    /// Extract a local .mcpb bundle instead of downloading. The
172    /// `STACKQL_MCP_BUNDLE` env var takes precedence over this.
173    pub fn bundle_path(mut self, path: impl Into<PathBuf>) -> Self {
174        self.acquisition.bundle_path = Some(path.into());
175        self
176    }
177
178    /// Embed the .mcpb bundle in your binary and extract it on first run:
179    /// `builder.bundle_bytes(include_bytes!("../stackql-mcp-linux-x64.mcpb"))`.
180    #[cfg(feature = "vendored")]
181    pub fn bundle_bytes(mut self, bytes: &'static [u8]) -> Self {
182        self.acquisition.bundle_bytes = Some(bytes);
183        self
184    }
185
186    /// Resolve the binary (acquiring it if needed) and return a
187    /// [`std::process::Command`] preloaded with the canonical launch args.
188    /// Blocking. The escape hatch for callers bringing their own MCP stack
189    /// or process supervision; stdio configuration is left to the caller.
190    pub fn command(&self) -> Result<std::process::Command> {
191        let binary = acquire::resolve_binary(&self.acquisition)?;
192        let approot = self.resolved_approot()?;
193        let mut cmd = std::process::Command::new(binary);
194        cmd.args(launch::launch_args(self.mode, &approot, self.auth.as_ref()));
195        Ok(cmd)
196    }
197
198    /// Acquire the binary if needed, spawn the server, and complete the MCP
199    /// handshake. Must be called from within a tokio runtime.
200    pub async fn start(self) -> Result<RunningServer> {
201        let approot = self.resolved_approot()?;
202        let acquisition = self.acquisition;
203        let binary = tokio::task::spawn_blocking(move || acquire::resolve_binary(&acquisition))
204            .await
205            .map_err(|e| Error::Mcp(format!("acquisition task failed: {e}")))??;
206
207        let mut child = tokio::process::Command::new(&binary)
208            .args(launch::launch_args(self.mode, &approot, self.auth.as_ref()))
209            .stdin(Stdio::piped())
210            .stdout(Stdio::piped())
211            // Diagnostics belong on stderr; let them flow through.
212            .stderr(Stdio::inherit())
213            .kill_on_drop(true)
214            .spawn()
215            .map_err(Error::Spawn)?;
216
217        let stdout = child
218            .stdout
219            .take()
220            .ok_or_else(|| Error::Mcp("child stdout not captured".into()))?;
221        let stdin = child
222            .stdin
223            .take()
224            .ok_or_else(|| Error::Mcp("child stdin not captured".into()))?;
225
226        let client = ()
227            .serve((stdout, stdin))
228            .await
229            .map_err(|e| Error::Mcp(format!("initialize failed: {e}")))?;
230
231        Ok(RunningServer {
232            child,
233            client,
234            binary,
235        })
236    }
237
238    fn resolved_approot(&self) -> Result<PathBuf> {
239        match &self.approot {
240            Some(p) => Ok(p.clone()),
241            None => cache::default_approot(),
242        }
243    }
244}
245
246/// A running embedded server: the child process handle plus a connected
247/// rmcp client. Derefs to the client, so rmcp peer methods
248/// (`list_all_tools`, `call_tool`, ...) are available directly.
249pub struct RunningServer {
250    child: tokio::process::Child,
251    client: RunningService<RoleClient, ()>,
252    binary: PathBuf,
253}
254
255impl RunningServer {
256    /// The connected rmcp client.
257    pub fn client(&self) -> &RunningService<RoleClient, ()> {
258        &self.client
259    }
260
261    /// OS process id of the server, if it is still running.
262    pub fn pid(&self) -> Option<u32> {
263        self.child.id()
264    }
265
266    /// Path of the stackql binary that was launched.
267    pub fn binary_path(&self) -> &std::path::Path {
268        &self.binary
269    }
270
271    /// Close the MCP session and stop the server process.
272    pub async fn shutdown(self) -> Result<()> {
273        let RunningServer {
274            mut child, client, ..
275        } = self;
276        // Cancelling drops the transport; the server sees EOF on stdin and
277        // exits. The kill is a backstop for a wedged process.
278        let _ = client.cancel().await;
279        let _ = child.kill().await;
280        Ok(())
281    }
282}
283
284impl Deref for RunningServer {
285    type Target = RunningService<RoleClient, ()>;
286
287    fn deref(&self) -> &Self::Target {
288        &self.client
289    }
290}