Skip to main content

vcon_server/
mcp.rs

1//! The MCP server: the HTTP API's operations as MCP tools for AI agents.
2//!
3//! It is served at `/mcp` on the TLS listener, in rmcp's stateless mode with
4//! JSON answers, behind the same authentication as every route. The stdio
5//! command in [`stdio`] forwards to it, so both transports pass one gate.
6
7mod admin;
8mod audit;
9pub mod budget;
10pub mod fence;
11mod gate;
12mod legal;
13pub mod stdio;
14mod vcons;
15
16use std::future::Future;
17use std::pin::Pin;
18use std::sync::Arc;
19use std::time::Instant;
20
21use axum::Router;
22use axum::body::Body;
23use axum::extract::{Request, State};
24use axum::middleware::Next;
25use axum::response::{IntoResponse, Response};
26use rmcp::handler::server::common::schema_for_type;
27use rmcp::model::{
28    CallToolRequestParams, CallToolResponse, CallToolResult, Extensions, Implementation,
29    JsonObject, ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool,
30    ToolAnnotations,
31};
32use rmcp::service::{RequestContext, RoleServer};
33use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
34use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
35use rmcp::{ErrorData, ServerHandler};
36use serde::de::DeserializeOwned;
37use serde_json::Value;
38
39pub use gate::scope;
40
41use crate::access::{self, Caller};
42use crate::api::{self, App, Problem};
43use crate::audit::{Action, Actor, Entry};
44use crate::ids::TenantId;
45
46/// What every client is told the server holds.
47const INSTRUCTIONS: &str = "vcon-server stores vCons, encrypted at rest, and \
48keeps an audit log of every action. Its tools act only within the tenant of \
49the API key that calls them, and only within that key's scopes; each call is \
50recorded in the audit log.";
51
52/// The limits from configuration that the MCP server applies.
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct Limits {
55    /// Tool calls one key may make in a one-second window.
56    pub calls_per_second: u32,
57    /// Tool calls in flight on the server.
58    pub concurrent: u32,
59    /// The largest result, in bytes.
60    pub response_bytes: usize,
61}
62
63impl Limits {
64    /// sipnab's defaults for the same limits.
65    pub const DEFAULT: Limits = Limits {
66        calls_per_second: 100,
67        concurrent: 100,
68        response_bytes: 1 << 20,
69    };
70}
71
72/// The hints a tool states besides `openWorldHint: false`.
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub(crate) struct Hints {
75    pub(crate) read_only: bool,
76    pub(crate) destructive: bool,
77    pub(crate) idempotent: bool,
78}
79
80impl Hints {
81    pub(crate) const NONE: Hints = Hints {
82        read_only: false,
83        destructive: false,
84        idempotent: false,
85    };
86}
87
88/// A tool named `name` taking `A` as its arguments, with `hints`.
89pub(crate) fn tool<A: rmcp::schemars::JsonSchema + 'static>(
90    name: &'static str,
91    description: &'static str,
92    hints: Hints,
93) -> Tool {
94    Tool::new(name, description, schema_for_type::<A>()).with_annotations(
95        ToolAnnotations::new()
96            .read_only(hints.read_only)
97            .destructive(hints.destructive)
98            .idempotent(hints.idempotent)
99            .open_world(false),
100    )
101}
102
103/// Why a tool did not answer with a result.
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105#[must_use]
106pub(crate) enum Outcome {
107    /// Its arguments do not fit its type, or a value in them names nothing
108    /// the tool can read; the gate records invalid-arguments.
109    Arguments,
110    /// The tool refused before the layer it calls saw the call; the gate
111    /// records the problem, and the result carries it.
112    Refused(Problem),
113    /// The layer the tool calls refused or failed, and recorded it; the
114    /// result carries the problem.
115    Answered(Problem),
116}
117
118impl std::fmt::Display for Outcome {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        match self {
121            Outcome::Arguments => f.write_str("the arguments do not fit the tool"),
122            Outcome::Refused(problem) => write!(f, "refused: {}", problem.name().1),
123            Outcome::Answered(problem) => write!(f, "answered: {}", problem.name().1),
124        }
125    }
126}
127
128impl std::error::Error for Outcome {}
129
130/// A tool's running call.
131pub(crate) type Running<'a> = Pin<Box<dyn Future<Output = Result<Value, Outcome>> + Send + 'a>>;
132
133/// What a tool does with its caller and arguments.
134pub(crate) type Run = for<'a> fn(&'a Server, Caller, JsonObject) -> Running<'a>;
135
136/// One registered tool: what clients are shown, the action it is recorded
137/// under, and what it runs.
138#[derive(Debug, Clone)]
139pub(crate) struct Spec {
140    /// The tool as `tools/list` shows it.
141    pub(crate) tool: Tool,
142    /// The action its calls are recorded under.
143    pub(crate) action: Action,
144    /// What a call runs.
145    pub(crate) run: Run,
146}
147
148/// A tool's arguments as its type, or the refusal when they do not fit.
149pub(crate) fn arguments<A: DeserializeOwned>(arguments: JsonObject) -> Result<A, Outcome> {
150    serde_json::from_value(Value::Object(arguments)).map_err(|_| Outcome::Arguments)
151}
152
153/// The MCP server's routes: `/mcp`, behind authentication and counting,
154/// with every tool.
155pub fn router(app: Arc<App>) -> Router {
156    router_with(app, tools())
157}
158
159/// Every tool, by area.
160pub(crate) fn tools() -> Vec<Spec> {
161    let mut all = vcons::specs();
162    all.extend(legal::specs());
163    all.extend(audit::specs());
164    all.extend(admin::specs());
165    all
166}
167
168/// The caller's Store and its caller for Store, for a tool needing `scope`;
169/// a tenant being removed is unavailable.
170pub(crate) struct Opened {
171    pub(crate) store: Arc<crate::store::Store>,
172    pub(crate) who: crate::store::Caller,
173}
174
175impl std::fmt::Debug for Opened {
176    // The store holds the tenant's keys; only the caller is shown.
177    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178        f.debug_struct("Opened")
179            .field("who", &self.who)
180            .finish_non_exhaustive()
181    }
182}
183
184pub(crate) async fn opened(
185    server: &Server,
186    caller: Caller,
187    scope: access::Scope,
188) -> Result<Opened, Outcome> {
189    let tenant =
190        access::tenant_for(&caller, scope).map_err(|_| Outcome::Refused(Problem::Forbidden))?;
191    let tenants = server.app.tenants.clone();
192    let store = api::blocking(move || tenants.store(tenant))
193        .await
194        .map_err(|_| Outcome::Refused(Problem::Unavailable))?;
195    Ok(Opened {
196        store,
197        who: crate::store::Caller { key: caller.key },
198    })
199}
200
201/// A stored version's JSON text, inflated within `limit` when it was
202/// stored gzipped, with the characters the fence never holds raw escaped.
203pub(crate) fn vcon_text(stored: &crate::store::Stored, limit: u64) -> Result<String, Outcome> {
204    use std::io::Read;
205    let damaged = || Outcome::Answered(Problem::Store(crate::store::StoreError::Damaged));
206    let text = if stored.gzip {
207        let mut inflated = Vec::new();
208        flate2::read::MultiGzDecoder::new(stored.body.as_slice())
209            .take(limit.saturating_add(1))
210            .read_to_end(&mut inflated)
211            .map_err(|_| damaged())?;
212        if inflated.len() as u64 > limit {
213            return Err(Outcome::Answered(Problem::TooLarge));
214        }
215        inflated
216    } else {
217        stored.body.clone()
218    };
219    let text = String::from_utf8(text).map_err(|_| damaged())?;
220    Ok(fence::escape_json(&text))
221}
222
223/// A result holding a piece of `text`: `head`'s members, the piece fenced,
224/// whether more follows and where, and the provenance sentence; the piece
225/// from `cursor`, a byte offset an earlier page gave, as long as `room`
226/// allows.
227pub(crate) fn text_result(
228    head: Value,
229    text: &str,
230    cursor: Option<&str>,
231    room: usize,
232) -> Result<Value, Outcome> {
233    let start = match cursor {
234        None => 0,
235        Some(cursor) => cursor
236            .parse::<usize>()
237            .ok()
238            .filter(|&at| at <= text.len() && text.is_char_boundary(at))
239            .ok_or(Outcome::Arguments)?,
240    };
241    let build = |piece: &str, next: Option<usize>| {
242        let mut result = head.clone();
243        result["vcon"] = fence::fence(piece).into();
244        result["truncated"] = next.is_some().into();
245        result["next_cursor"] = next.map(|at| at.to_string()).into();
246        result["provenance"] = fence::PROVENANCE.into();
247        result
248    };
249    let (piece, next) = budget::piece(text, start, |piece, more| {
250        // A following page's cursor is at most the text's length.
251        budget::size(&build(piece, more.then_some(text.len()))) <= room
252    });
253    Ok(build(piece, next))
254}
255
256/// The MCP server's routes with the tools in `specs`.
257pub(crate) fn router_with(app: Arc<App>, specs: Vec<Spec>) -> Router {
258    let server = Server {
259        app: app.clone(),
260        specs: Arc::new(specs),
261        gate: Arc::new(gate::Gate::new(app.mcp)),
262    };
263    let config = StreamableHttpServerConfig::default()
264        .with_legacy_session_mode(false)
265        .with_json_response(true)
266        .with_sse_keep_alive(None)
267        // TLS and the bearer key already stop a page that rebinds a name to
268        // this server, and a deployment's host name is its own.
269        .disable_allowed_hosts()
270        // A web page sends an Origin header; no page may call /mcp.
271        .enforce_origin_validation()
272        .with_max_request_body_bytes(app.limits.body);
273    let service: StreamableHttpService<Server, LocalSessionManager> =
274        StreamableHttpService::new(move || Ok(server.clone()), Default::default(), config);
275    // POST only: GET and DELETE answer 405 before any key is asked for, and
276    // the layer wraps this one route, never the router's fallback.
277    Router::new().route(
278        "/mcp",
279        axum::routing::post_service(service)
280            .route_layer(axum::middleware::from_fn_with_state(app, admit)),
281    )
282}
283
284/// The MCP request handler: the tools, and the gate every call passes.
285#[derive(Clone)]
286pub(crate) struct Server {
287    pub(crate) app: Arc<App>,
288    specs: Arc<Vec<Spec>>,
289    gate: Arc<gate::Gate>,
290}
291
292impl std::fmt::Debug for Server {
293    // The app holds keys and stores; the tools and the gate are shown.
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        let tools: Vec<&str> = self.specs.iter().map(|spec| &*spec.tool.name).collect();
296        f.debug_struct("Server")
297            .field("tools", &tools)
298            .field("gate", &self.gate)
299            .finish_non_exhaustive()
300    }
301}
302
303/// The caller a request to `/mcp` was admitted with.
304fn caller_of(extensions: &Extensions) -> Option<Caller> {
305    extensions
306        .get::<axum::http::request::Parts>()
307        .and_then(|parts| parts.extensions.get::<Caller>())
308        .copied()
309}
310
311/// Whether `caller` may call the tool, by the scope rule.
312fn allows(caller: &Caller, tool: &Tool) -> bool {
313    scope(&tool.name, tool.annotations.as_ref())
314        .is_some_and(|scope| access::require(caller, scope).is_ok())
315}
316
317impl Server {
318    /// Records a refusal in the caller's log and returns it, or audit-failed
319    /// when the record cannot be written.
320    async fn recorded(&self, caller: Caller, action: Action, problem: Problem) -> ErrorData {
321        let entry = refusal(caller, action, problem);
322        gate::refused(self.app.refuse(caller.tenant, entry, problem).await, None)
323    }
324
325    /// Counts a refusal in the caller's log and returns it, or audit-failed
326    /// when the count cannot be written.
327    async fn counted(
328        &self,
329        caller: Caller,
330        action: Action,
331        problem: Problem,
332        retry: Option<std::time::Duration>,
333    ) -> ErrorData {
334        if count(&self.app, caller.tenant, refusal(caller, action, problem)).await {
335            gate::refused(problem, retry)
336        } else {
337            gate::refused(Problem::AuditFailed, None)
338        }
339    }
340}
341
342/// The record of a call the gate refused.
343fn refusal(caller: Caller, action: Action, problem: Problem) -> Entry {
344    Entry {
345        actor: Actor::Key(caller.key),
346        action,
347        target: None,
348        version: None,
349        status: problem.status().as_u16(),
350    }
351}
352
353impl ServerHandler for Server {
354    fn get_info(&self) -> ServerConfig {
355        let mut info = ServerConfig::new(ServerCapabilities::builder().enable_tools().build());
356        // rmcp's default names rmcp itself; the server names itself.
357        info.server_info = Implementation::new("vcon-server", env!("CARGO_PKG_VERSION"));
358        info.instructions = Some(INSTRUCTIONS.to_owned());
359        info
360    }
361
362    /// The tools the caller's key may call, by the scope rule.
363    async fn list_tools(
364        &self,
365        _request: Option<PaginatedRequestParams>,
366        context: RequestContext<RoleServer>,
367    ) -> Result<ListToolsResult, ErrorData> {
368        let caller = caller_of(&context.extensions);
369        let mut tools: Vec<Tool> = self
370            .specs
371            .iter()
372            .filter(|spec| caller.is_some_and(|caller| allows(&caller, &spec.tool)))
373            .map(|spec| spec.tool.clone())
374            .collect();
375        tools.sort_by(|a, b| a.name.cmp(&b.name));
376        Ok(ListToolsResult::with_all_items(tools))
377    }
378
379    /// The gate every tool call passes, in order: the key's rate, a free
380    /// slot, a known tool, the tool's scope, then the tool with arguments
381    /// that fit it. Each refusal is recorded, or counted for rate and
382    /// capacity, before it is sent.
383    async fn call_tool(
384        &self,
385        request: CallToolRequestParams,
386        context: RequestContext<RoleServer>,
387    ) -> Result<CallToolResponse, ErrorData> {
388        let Some(caller) = caller_of(&context.extensions) else {
389            // Every request reaches here through admit, which sets it.
390            return Err(gate::refused(Problem::Unauthorized, None));
391        };
392        let spec = self
393            .specs
394            .iter()
395            .find(|spec| spec.tool.name == request.name);
396        let action = spec.map_or(Action::Mcp, |spec| spec.action);
397        let limit = self.gate.limits.calls_per_second;
398        if let Err(retry) = self.gate.windows.take(caller.key, limit, Instant::now()) {
399            return Err(self
400                .counted(caller, action, Problem::RateLimited, Some(retry))
401                .await);
402        }
403        let Some(_slot) = self.gate.slot() else {
404            return Err(self
405                .counted(caller, action, Problem::AtCapacity, None)
406                .await);
407        };
408        let Some(spec) = spec else {
409            return Err(self.recorded(caller, action, Problem::UnknownTool).await);
410        };
411        if !allows(&caller, &spec.tool) {
412            return Err(self.recorded(caller, action, Problem::Forbidden).await);
413        }
414        // The tool runs in its own task, so a panic in it is answered and
415        // recorded as internal instead of leaving the request unanswered.
416        let (run, server, arguments) = (
417            spec.run,
418            self.clone(),
419            request.arguments.unwrap_or_default(),
420        );
421        let ran = tokio::spawn(async move { (run)(&server, caller, arguments).await }).await;
422        let Ok(ran) = ran else {
423            let failed = Problem::Store(crate::store::StoreError::Io);
424            let entry = refusal(caller, action, failed);
425            let problem = self.app.refuse(caller.tenant, entry, failed).await;
426            return Ok(CallToolResult::structured_error(problem.document()).into());
427        };
428        match ran {
429            Ok(value) => Ok(CallToolResult::structured(value).into()),
430            Err(Outcome::Arguments) => Err(self
431                .recorded(caller, action, Problem::InvalidArguments)
432                .await),
433            Err(Outcome::Refused(problem)) => {
434                let entry = refusal(caller, action, problem);
435                let problem = self.app.refuse(caller.tenant, entry, problem).await;
436                Ok(CallToolResult::structured_error(problem.document()).into())
437            }
438            Err(Outcome::Answered(problem)) => {
439                Ok(CallToolResult::structured_error(problem.document()).into())
440            }
441        }
442    }
443}
444
445/// Counts `entry` in `tenant`'s log, or in the admin log for an operator key
446/// or a tenant whose log cannot be opened; false when it cannot be.
447async fn count(app: &Arc<App>, tenant: Option<TenantId>, entry: Entry) -> bool {
448    let tenants = app.tenants.clone();
449    api::blocking(move || match tenant.and_then(|t| tenants.store(t).ok()) {
450        Some(store) => store.count(&entry).is_ok(),
451        None => tenants.count_admin(&entry).is_ok(),
452    })
453    .await
454}
455
456/// Every request to `/mcp`, before rmcp reads it: the bearer key is
457/// authenticated as on every route, the caller's tenant opened, the body
458/// read within the limits, and the caller passed on inside the request.
459/// Each answered request is then counted in the key's log; one that cannot
460/// be counted answers audit-failed.
461async fn admit(State(app): State<Arc<App>>, request: Request, next: Next) -> Response {
462    let (mut parts, body) = request.into_parts();
463    let caller = match api::authenticated(&app, &parts.headers).await {
464        Ok(caller) => caller,
465        Err(refused) => return refused.into_response(),
466    };
467    // A tenant being removed answers 503, as on every route, and is counted
468    // in the admin log, since its own may already be gone.
469    let tenants = app.tenants.clone();
470    let closing = caller
471        .tenant
472        .is_some_and(|tenant| tenants.store(tenant).is_err());
473    let answer = if closing {
474        Problem::Unavailable.into_response()
475    } else {
476        match api::read_body(body, &parts.headers, &app.limits).await {
477            Ok(body) => {
478                parts.extensions.insert(caller);
479                next.run(Request::from_parts(parts, Body::from(body))).await
480            }
481            Err(refused) => refused.into_response(),
482        }
483    };
484    let entry = Entry {
485        actor: Actor::Key(caller.key),
486        action: Action::Mcp,
487        target: None,
488        version: None,
489        status: answer.status().as_u16(),
490    };
491    if count(&app, caller.tenant, entry).await {
492        answer
493    } else {
494        Problem::AuditFailed.into_response()
495    }
496}
497
498#[cfg(test)]
499mod tests;