1mod 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
46const 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct Limits {
55 pub calls_per_second: u32,
57 pub concurrent: u32,
59 pub response_bytes: usize,
61}
62
63impl Limits {
64 pub const DEFAULT: Limits = Limits {
66 calls_per_second: 100,
67 concurrent: 100,
68 response_bytes: 1 << 20,
69 };
70}
71
72#[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
88pub(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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105#[must_use]
106pub(crate) enum Outcome {
107 Arguments,
110 Refused(Problem),
113 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
130pub(crate) type Running<'a> = Pin<Box<dyn Future<Output = Result<Value, Outcome>> + Send + 'a>>;
132
133pub(crate) type Run = for<'a> fn(&'a Server, Caller, JsonObject) -> Running<'a>;
135
136#[derive(Debug, Clone)]
139pub(crate) struct Spec {
140 pub(crate) tool: Tool,
142 pub(crate) action: Action,
144 pub(crate) run: Run,
146}
147
148pub(crate) fn arguments<A: DeserializeOwned>(arguments: JsonObject) -> Result<A, Outcome> {
150 serde_json::from_value(Value::Object(arguments)).map_err(|_| Outcome::Arguments)
151}
152
153pub fn router(app: Arc<App>) -> Router {
156 router_with(app, tools())
157}
158
159pub(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
168pub(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 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
201pub(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
223pub(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 budget::size(&build(piece, more.then_some(text.len()))) <= room
252 });
253 Ok(build(piece, next))
254}
255
256pub(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 .disable_allowed_hosts()
270 .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 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#[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 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
303fn 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
311fn 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 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 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
342fn 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 info.server_info = Implementation::new("vcon-server", env!("CARGO_PKG_VERSION"));
358 info.instructions = Some(INSTRUCTIONS.to_owned());
359 info
360 }
361
362 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 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 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 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
445async 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
456async 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 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;