#![cfg(feature = "mcp")]
use std::collections::{HashMap, VecDeque};
use std::convert::Infallible;
use std::fmt::Write as _;
use std::pin::Pin;
use std::sync::Arc;
use axum::body::{Body, Bytes};
use axum::http::{HeaderMap, HeaderName, HeaderValue, StatusCode, header};
use axum::response::{IntoResponse, Response};
use futures::{FutureExt as _, Stream, StreamExt as _};
use serde_json::{Value, json};
use tower::ServiceExt as _;
use crate::sse::{Event, Sse};
use crate::openapi::{ApiDoc, SchemaDirection, schema_entry_to_value_in};
mod catalog;
mod modern;
mod tasks;
pub use catalog::McpCatalog;
pub use tasks::McpTaskConfig;
const DEFAULT_PROTOCOL_VERSION: &str = "2025-06-18";
const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &["2025-06-18", "2025-03-26", "2024-11-05"];
const MCP_REQUEST_HEADERS: &[&str] = &[
"mcp-protocol-version",
"mcp-session-id",
"mcp-name",
"mcp-method",
"payment-signature",
];
const MAX_TOOL_RESPONSE_BYTES: usize = 10 * 1024 * 1024;
const FORWARDED_HEADERS: &[&str] = &[
"authorization",
"cookie",
"idempotency-key",
"host",
"forwarded",
"x-forwarded-for",
"x-forwarded-host",
"x-forwarded-proto",
"x-real-ip",
"accept-language",
"payment-signature",
];
const REPLAYED_RESPONSE_HEADERS: &[&str] = &["set-cookie", "payment-required", "payment-response"];
pub(crate) type McpEndpointLayer = Box<
dyn FnOnce(axum::Router<crate::state::AppState>) -> axum::Router<crate::state::AppState> + Send,
>;
pub struct McpRuntime {
pub mount_path: String,
pub expose_all: bool,
pub(crate) endpoint_layer: Option<McpEndpointLayer>,
pub(crate) tasks: McpTaskConfig,
pub(crate) instructions: Option<String>,
pub(crate) catalog: McpCatalog,
}
impl McpRuntime {
#[must_use]
pub fn new(mount_path: impl Into<String>) -> Self {
Self {
mount_path: mount_path.into(),
expose_all: false,
endpoint_layer: None,
tasks: McpTaskConfig::default(),
instructions: None,
catalog: McpCatalog::Full,
}
}
}
#[derive(Clone, Debug)]
#[allow(clippy::struct_excessive_bools)] struct McpTool {
name: String,
description: Option<String>,
input_schema: Value,
annotations: Value,
method: String,
path_template: String,
path_params: Vec<String>,
has_body: bool,
has_query: bool,
streams: bool,
empty_body: bool,
task: bool,
agent_authority: Option<&'static crate::agent_authority::AgentAuthority>,
core: bool,
tags: Vec<String>,
}
impl McpTool {
fn descriptor(&self) -> Value {
let mut obj = serde_json::Map::new();
obj.insert("name".into(), json!(self.name));
if let Some(desc) = &self.description {
obj.insert("description".into(), json!(desc));
}
obj.insert("inputSchema".into(), self.input_schema.clone());
obj.insert("annotations".into(), self.annotations.clone());
if self.core {
obj.insert("_meta".into(), json!({ catalog::ALWAYS_LOAD: true }));
}
Value::Object(obj)
}
}
pub(crate) struct McpWiring {
pub cors: crate::config::CorsConfig,
pub trusted_hosts: crate::router::TrustedHostPolicy,
pub tenant_header: Option<String>,
pub csrf_header: String,
pub envelope_rate_limited: bool,
pub envelope_load_shed: bool,
pub state: crate::state::AppState,
pub session_cookie: String,
pub tasks: McpTaskConfig,
pub instructions: Option<String>,
pub catalog: McpCatalog,
}
pub struct McpServer {
tools: Vec<McpTool>,
by_name: HashMap<String, usize>,
dispatch: axum::Router,
cors: crate::config::CorsConfig,
trusted_hosts: crate::router::TrustedHostPolicy,
tenant_header: Option<String>,
csrf_header: String,
envelope_rate_limited: bool,
envelope_load_shed: bool,
state: crate::state::AppState,
session_cookie: String,
tasks: tasks::TaskStore,
has_task_tools: bool,
server_name: String,
server_version: String,
instructions: Option<String>,
catalog: McpCatalog,
meta_tools: Vec<Value>,
search_index: Vec<catalog::ToolWords>,
rate_limiter: crate::agent_authority::rate::RateLimiter,
}
impl McpServer {
fn origin_allowed(&self, origin: &str, host: Option<&str>, scheme: Option<&str>) -> bool {
if let Some(host) = host
&& is_same_origin(origin, host, scheme)
&& crate::router::extract_host_without_port(host)
.is_some_and(|h| self.trusted_hosts.allows_host(&h.to_ascii_lowercase()))
{
return true;
}
self.cors
.allowed_origins
.iter()
.any(|allowed| allowed == "*" || allowed == origin)
}
}
fn is_same_origin(origin: &str, host: &str, scheme: Option<&str>) -> bool {
let Some((origin_scheme, origin_authority)) = origin.split_once("://") else {
return false;
};
if scheme.is_some_and(|s| !s.eq_ignore_ascii_case(origin_scheme)) {
return false;
}
let host_scheme = scheme.unwrap_or(origin_scheme);
authority_matches(origin_authority, origin_scheme, host, host_scheme)
}
fn authority_matches(a: &str, a_scheme: &str, b: &str, b_scheme: &str) -> bool {
let (a_host, a_port) = split_host_port(a);
let (b_host, b_port) = split_host_port(b);
if !a_host.eq_ignore_ascii_case(b_host) {
return false;
}
a_port.or_else(|| default_port(a_scheme)) == b_port.or_else(|| default_port(b_scheme))
}
fn split_host_port(authority: &str) -> (&str, Option<&str>) {
if authority.starts_with('[') {
if let Some(close) = authority.find(']') {
let host = &authority[..=close];
let port = authority[close + 1..]
.strip_prefix(':')
.filter(|p| !p.is_empty());
return (host, port);
}
return (authority, None);
}
match authority.rsplit_once(':') {
Some((host, port)) if !port.is_empty() && port.bytes().all(|c| c.is_ascii_digit()) => {
(host, Some(port))
}
_ => (authority, None),
}
}
fn default_port(scheme: &str) -> Option<&'static str> {
match scheme.to_ascii_lowercase().as_str() {
"https" => Some("443"),
"http" => Some("80"),
_ => None,
}
}
pub(crate) fn should_expose(doc: &ApiDoc, expose_all: bool) -> bool {
if doc.hidden || doc.mcp_exclude {
return false;
}
if doc.mcp_stream {
if doc.mcp_tool {
return true;
}
return expose_all && is_read_only(doc.method);
}
if doc.response.is_none() && !has_empty_body_contract(doc.success_status) {
return false;
}
if doc.mcp_tool {
return true;
}
if expose_all {
return is_read_only(doc.method);
}
false
}
fn is_read_only(method: &str) -> bool {
matches!(method.to_ascii_uppercase().as_str(), "GET" | "HEAD")
}
const fn has_empty_body_contract(status: u16) -> bool {
matches!(status, 204 | 205)
}
fn annotations_for(
method: &str,
title: &str,
authority: Option<&'static crate::agent_authority::AgentAuthority>,
) -> Value {
use crate::agent_authority::Reversibility;
let upper = method.to_ascii_uppercase();
let read_only = is_read_only(&upper);
let mut obj = serde_json::Map::new();
obj.insert("title".into(), json!(title));
obj.insert("readOnlyHint".into(), json!(read_only));
let verb_is_destructive = upper == "DELETE";
if let Some(authority) = authority {
let destructive =
authority.grant.reversibility != Reversibility::Reversible || verb_is_destructive;
obj.insert("destructiveHint".into(), json!(destructive));
} else if verb_is_destructive {
obj.insert("destructiveHint".into(), json!(true));
}
Value::Object(obj)
}
fn build_input_schema(
doc: &ApiDoc,
components: &serde_json::Map<String, Value>,
index: &crate::openapi::SchemaComponentIndex,
) -> Value {
let mut properties = serde_json::Map::new();
let mut required: Vec<Value> = Vec::new();
let mut defs = serde_json::Map::new();
for param in doc.path_params {
let name = param.strip_prefix('*').unwrap_or(param);
properties.insert(name.to_owned(), json!({ "type": "string" }));
required.push(json!(name));
}
if let Some(query) = &doc.query_schema {
let schema = rewrite_refs(
schema_entry_to_value_in(query, index, SchemaDirection::Request),
components,
&mut defs,
);
properties.insert("query".to_owned(), schema);
}
if let Some(body) = &doc.request_body {
let schema = rewrite_refs(
schema_entry_to_value_in(body, index, SchemaDirection::Request),
components,
&mut defs,
);
properties.insert("body".to_owned(), schema);
required.push(json!("body"));
}
let mut schema = serde_json::Map::new();
schema.insert("type".into(), json!("object"));
schema.insert("properties".into(), Value::Object(properties));
if !required.is_empty() {
schema.insert("required".into(), Value::Array(required));
}
if !defs.is_empty() {
schema.insert("$defs".into(), Value::Object(defs));
}
Value::Object(schema)
}
fn rewrite_refs(
value: Value,
components: &serde_json::Map<String, Value>,
defs: &mut serde_json::Map<String, Value>,
) -> Value {
match value {
Value::Object(map) => {
if let Some(Value::String(reference)) = map.get("$ref")
&& let Some(name) = reference.strip_prefix("#/components/schemas/")
{
let name = name.to_owned();
let local = format!("#/$defs/{name}");
if !defs.contains_key(&name) {
defs.insert(name.clone(), Value::Null);
let resolved = components
.get(&name)
.cloned()
.unwrap_or_else(|| json!({ "type": "object", "title": name.clone() }));
let resolved = rewrite_refs(resolved, components, defs);
defs.insert(name, resolved);
}
return json!({ "$ref": local });
}
let rewritten: serde_json::Map<String, Value> = map
.into_iter()
.map(|(k, v)| (k, rewrite_refs(v, components, defs)))
.collect();
Value::Object(rewritten)
}
Value::Array(items) => Value::Array(
items
.into_iter()
.map(|v| rewrite_refs(v, components, defs))
.collect(),
),
other => other,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SchemaDegradation {
OpaqueQuery,
OpaqueBody,
}
fn resolve_local_ref<'a>(root: &'a Value, node: &'a Value) -> &'a Value {
if let Some(name) = node
.get("$ref")
.and_then(Value::as_str)
.and_then(|r| r.strip_prefix("#/$defs/"))
&& let Some(def) = root.get("$defs").and_then(|defs| defs.get(name))
{
return def;
}
node
}
fn detect_schema_degradations(
input_schema: &Value,
has_query: bool,
has_body: bool,
) -> Vec<SchemaDegradation> {
let mut out = Vec::new();
let props = input_schema.get("properties");
if has_body
&& let Some(body) = props.and_then(|p| p.get("body"))
&& crate::openapi::is_opaque_object_schema(resolve_local_ref(input_schema, body))
{
out.push(SchemaDegradation::OpaqueBody);
}
if has_query
&& let Some(query) = props.and_then(|p| p.get("query"))
&& crate::openapi::is_opaque_object_schema(resolve_local_ref(input_schema, query))
{
out.push(SchemaDegradation::OpaqueQuery);
}
out
}
fn warn_on_degraded_input_schema(doc: &ApiDoc, input_schema: &Value) {
for degradation in detect_schema_degradations(
input_schema,
doc.query_schema.is_some(),
doc.request_body.is_some(),
) {
match degradation {
SchemaDegradation::OpaqueBody => tracing::warn!(
operation_id = doc.operation_id,
method = doc.method,
path = doc.path,
"MCP tool body has no field-level schema (opaque object); derive or \
impl `OpenApiSchema` on the `Json<T>` body type so the tool advertises \
its real fields instead of a bare `{{\"type\":\"object\"}}` placeholder"
),
SchemaDegradation::OpaqueQuery => tracing::warn!(
operation_id = doc.operation_id,
method = doc.method,
path = doc.path,
"MCP tool query has no field-level schema (opaque object); derive or \
impl `OpenApiSchema` on the `Query<T>` type so the tool advertises \
its real query parameters instead of a bare `{{\"type\":\"object\"}}` placeholder"
),
}
}
}
#[must_use]
pub fn derive_tools(
docs: &[ApiDoc],
expose_all: bool,
openapi: Option<&crate::openapi::OpenApiConfig>,
) -> Vec<McpToolInfo> {
let refs: Vec<&ApiDoc> = docs.iter().collect();
let config = openapi.cloned().unwrap_or_else(|| {
crate::openapi::OpenApiConfig::new("autumn-mcp", env!("CARGO_PKG_VERSION"))
});
let spec = crate::openapi::generate_spec(&config, &refs);
let components = spec
.components
.as_ref()
.map(|c| serde_json::to_value(&c.schemas).unwrap_or(Value::Null))
.and_then(|v| v.as_object().cloned())
.unwrap_or_default();
let index = crate::openapi::schema_component_index_for(&config, &refs);
let mut tools = Vec::new();
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
for doc in docs {
if (doc.mcp_tool || (expose_all && is_read_only(doc.method)))
&& doc.response.is_none()
&& !has_empty_body_contract(doc.success_status)
&& !doc.mcp_stream
&& !doc.mcp_exclude
&& !doc.hidden
{
tracing::warn!(
operation_id = doc.operation_id,
method = doc.method,
path = doc.path,
"skipping MCP exposure: endpoint has no JSON response schema; \
eligible tools return Json<T>, declare an empty-body status \
(204/205), or opt in as streaming (`stream`/`Route::mcp_stream`) \
— HTML/Maud routes are not eligible"
);
continue;
}
if !should_expose(doc, expose_all) {
continue;
}
if !seen.insert(doc.operation_id) {
tracing::warn!(
operation_id = doc.operation_id,
method = doc.method,
path = doc.path,
"duplicate MCP tool name; keeping the first registration and \
skipping this duplicate (set a distinct operation_id to expose both)"
);
continue;
}
let title = doc.summary.unwrap_or(doc.operation_id);
let input_schema = build_input_schema(doc, &components, &index);
warn_on_degraded_input_schema(doc, &input_schema);
tools.push(McpToolInfo {
name: doc.operation_id.to_owned(),
description: doc.description.or(doc.summary).map(str::to_owned),
input_schema,
annotations: annotations_for(doc.method, title, doc.agent_authority),
method: doc.method.to_owned(),
path_template: doc.path.to_owned(),
path_params: doc.path_params.iter().map(|p| (*p).to_owned()).collect(),
has_body: doc.request_body.is_some(),
has_query: doc.query_schema.is_some(),
streams: doc.mcp_stream,
empty_body: doc.response.is_none()
&& has_empty_body_contract(doc.success_status)
&& !doc.mcp_stream,
task: doc.mcp_task,
agent_authority: doc.agent_authority,
core: doc.mcp_core,
tags: doc.tags.iter().map(|tag| (*tag).to_owned()).collect(),
});
}
tools
}
pub(crate) fn drop_reserved_tools(tools: Vec<McpToolInfo>, mode: McpCatalog) -> Vec<McpToolInfo> {
catalog::drop_reserved(tools, mode)
}
#[derive(Clone, Debug)]
#[allow(clippy::struct_excessive_bools)] pub struct McpToolInfo {
name: String,
description: Option<String>,
input_schema: Value,
annotations: Value,
method: String,
path_template: String,
path_params: Vec<String>,
has_body: bool,
has_query: bool,
streams: bool,
empty_body: bool,
task: bool,
agent_authority: Option<&'static crate::agent_authority::AgentAuthority>,
core: bool,
tags: Vec<String>,
}
impl McpToolInfo {
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn description(&self) -> Option<&str> {
self.description.as_deref()
}
#[must_use]
pub const fn input_schema(&self) -> &Value {
&self.input_schema
}
#[must_use]
pub const fn annotations(&self) -> &Value {
&self.annotations
}
#[must_use]
pub fn method(&self) -> &str {
&self.method
}
#[must_use]
pub fn path_template(&self) -> &str {
&self.path_template
}
#[must_use]
pub fn path_params(&self) -> &[String] {
&self.path_params
}
#[must_use]
pub const fn has_body(&self) -> bool {
self.has_body
}
#[must_use]
pub const fn has_query(&self) -> bool {
self.has_query
}
#[must_use]
pub const fn streams(&self) -> bool {
self.streams
}
#[must_use]
pub const fn supports_tasks(&self) -> bool {
self.task
}
#[must_use]
pub const fn is_core(&self) -> bool {
self.core
}
#[must_use]
pub const fn agent_authority(&self) -> Option<&'static crate::agent_authority::AgentAuthority> {
self.agent_authority
}
}
impl McpServer {
#[must_use]
pub(crate) fn new(tools: Vec<McpToolInfo>, dispatch: axum::Router, wiring: McpWiring) -> Self {
let progressive = wiring.catalog == McpCatalog::Progressive;
let tools: Vec<McpTool> = catalog::drop_reserved(tools, wiring.catalog)
.into_iter()
.map(|t| McpTool {
name: t.name,
description: t.description,
input_schema: t.input_schema,
annotations: t.annotations,
method: t.method,
path_template: t.path_template,
path_params: t.path_params,
has_body: t.has_body,
has_query: t.has_query,
streams: t.streams,
empty_body: t.empty_body,
task: t.task,
agent_authority: t.agent_authority,
core: t.core,
tags: t.tags,
})
.collect();
let (meta_tools, search_index) = if progressive {
(catalog::meta_descriptors(&tools), catalog::index(&tools))
} else {
(Vec::new(), Vec::new())
};
let has_task_tools = tools.iter().any(|t| t.task);
let by_name = tools
.iter()
.enumerate()
.map(|(i, t)| (t.name.clone(), i))
.collect();
Self {
tools,
by_name,
dispatch,
cors: wiring.cors,
trusted_hosts: wiring.trusted_hosts,
tenant_header: wiring.tenant_header,
csrf_header: wiring.csrf_header,
envelope_rate_limited: wiring.envelope_rate_limited,
envelope_load_shed: wiring.envelope_load_shed,
state: wiring.state,
session_cookie: wiring.session_cookie,
tasks: tasks::TaskStore::new(wiring.tasks),
has_task_tools,
server_name: "autumn-mcp".to_owned(),
server_version: env!("CARGO_PKG_VERSION").to_owned(),
instructions: catalog::instructions(wiring.instructions.as_deref(), wiring.catalog),
catalog: wiring.catalog,
meta_tools,
search_index,
rate_limiter: crate::agent_authority::rate::RateLimiter::default(),
}
}
}
pub(crate) fn build_mcp_router(
mount_path: &str,
tools: Vec<McpToolInfo>,
dispatch: axum::Router,
wiring: McpWiring,
endpoint_layer: Option<McpEndpointLayer>,
) -> axum::Router<crate::state::AppState> {
let server = Arc::new(McpServer::new(tools, dispatch, wiring));
tracing::debug!(
path = mount_path,
tools = server.tools.len(),
"Mounted MCP endpoint"
);
let mut rpc = axum::Router::<crate::state::AppState>::new()
.route(
mount_path,
axum::routing::get(serve_mcp_get).post(serve_mcp),
)
.layer(axum::extract::Extension(Arc::clone(&server)));
if let Some(layer_fn) = endpoint_layer {
rpc = layer_fn(rpc);
}
let guard_server = Arc::clone(&server);
rpc = rpc.layer(axum::middleware::from_fn(move |req, next| {
mcp_host_origin_guard(Arc::clone(&guard_server), req, next)
}));
let preflight = axum::Router::<crate::state::AppState>::new()
.route(mount_path, axum::routing::options(serve_mcp_options))
.layer(axum::extract::Extension(server));
rpc.merge(preflight)
}
pub(crate) fn apply_mcp_cors_layer(
router: axum::Router<crate::state::AppState>,
cors: &crate::config::CorsConfig,
) -> axum::Router<crate::state::AppState> {
let cors = cors.clone();
router.layer(axum::middleware::from_fn(
move |req: axum::extract::Request, next: axum::middleware::Next| {
let cors = cors.clone();
async move {
let origin = req
.headers()
.get(header::ORIGIN)
.and_then(|v| v.to_str().ok())
.map(str::to_owned);
let mut response = next.run(req).await;
apply_cors_headers(&cors, origin.as_deref(), &mut response);
response
}
},
))
}
async fn serve_mcp_options(
axum::extract::Extension(server): axum::extract::Extension<Arc<McpServer>>,
headers: HeaderMap,
) -> Response {
use axum::http::HeaderValue;
let cors = &server.cors;
let mut out = HeaderMap::new();
out.insert(header::VARY, HeaderValue::from_static("origin"));
let origin = headers.get(header::ORIGIN).and_then(|o| o.to_str().ok());
let Some(allow_origin) = cors_allow_origin(cors, origin) else {
out.insert(
header::ALLOW,
HeaderValue::from_static("GET, POST, OPTIONS"),
);
return (StatusCode::NO_CONTENT, out).into_response();
};
if let Ok(v) = HeaderValue::from_str(&allow_origin) {
out.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, v);
}
if let Ok(v) = HeaderValue::from_str(&cors.allowed_methods.join(", ")) {
out.insert(header::ACCESS_CONTROL_ALLOW_METHODS, v);
}
let mut allow_headers = cors.allowed_headers.clone();
for extra in MCP_REQUEST_HEADERS {
if !allow_headers.iter().any(|h| h.eq_ignore_ascii_case(extra)) {
allow_headers.push((*extra).to_owned());
}
}
if let Ok(v) = HeaderValue::from_str(&allow_headers.join(", ")) {
out.insert(header::ACCESS_CONTROL_ALLOW_HEADERS, v);
}
if let Ok(v) = HeaderValue::from_str(&cors.max_age_secs.to_string()) {
out.insert(header::ACCESS_CONTROL_MAX_AGE, v);
}
if cors.allow_credentials {
out.insert(
header::ACCESS_CONTROL_ALLOW_CREDENTIALS,
HeaderValue::from_static("true"),
);
}
(StatusCode::NO_CONTENT, out).into_response()
}
fn cors_allow_origin(cors: &crate::config::CorsConfig, origin: Option<&str>) -> Option<String> {
let origin = origin?;
let allow_any = cors.allowed_origins.iter().any(|a| a == "*");
if !(allow_any || cors.allowed_origins.iter().any(|a| a == origin)) {
return None;
}
Some(if allow_any && !cors.allow_credentials {
"*".to_owned()
} else {
origin.to_owned()
})
}
fn apply_cors_headers(
cors: &crate::config::CorsConfig,
origin: Option<&str>,
response: &mut Response,
) {
use axum::http::HeaderValue;
let headers = response.headers_mut();
headers.insert(header::VARY, HeaderValue::from_static("origin"));
if let Some(allow_origin) = cors_allow_origin(cors, origin)
&& let Ok(v) = HeaderValue::from_str(&allow_origin)
{
headers.insert(header::ACCESS_CONTROL_ALLOW_ORIGIN, v);
if cors.allow_credentials {
headers.insert(
header::ACCESS_CONTROL_ALLOW_CREDENTIALS,
HeaderValue::from_static("true"),
);
}
}
}
async fn serve_mcp_get() -> Response {
(
StatusCode::METHOD_NOT_ALLOWED,
[(header::ALLOW, "POST")],
"MCP server-initiated streaming is not supported (POST JSON-RPC only)",
)
.into_response()
}
static SERVER_TIMING: HeaderName = HeaderName::from_static("server-timing");
struct ReplayContext<'a> {
headers: &'a HeaderMap,
identity: Option<&'a crate::security::ResolvedClientIdentity>,
peer: Option<std::net::SocketAddr>,
envelope_admission: Option<crate::middleware::EnvelopeAdmission>,
#[cfg(feature = "tls")]
client_cert: Option<std::sync::Arc<crate::tls::client_auth::ClientIdentity>>,
modern: bool,
}
async fn mcp_host_origin_guard(
server: Arc<McpServer>,
req: axum::extract::Request,
next: axum::middleware::Next,
) -> Response {
let identity = req
.extensions()
.get::<crate::security::ResolvedClientIdentity>();
let host = identity
.and_then(|id| id.host.as_deref())
.or_else(|| req.uri().authority().map(http::uri::Authority::as_str))
.or_else(|| {
req.headers()
.get(header::HOST)
.and_then(|h| h.to_str().ok())
});
let host_trusted = host
.and_then(crate::router::extract_host_without_port)
.map(|h| h.trim_end_matches('.').to_ascii_lowercase())
.filter(|h| !h.is_empty())
.map_or_else(
|| server.trusted_hosts.allows_missing_host(),
|h| server.trusted_hosts.allows_host(&h),
);
if !host_trusted {
return (StatusCode::BAD_REQUEST, "Invalid Host header").into_response();
}
if let Some(origin) = req.headers().get(header::ORIGIN) {
let origin = origin.to_str().unwrap_or("");
let scheme = identity.and_then(|id| id.scheme.as_deref());
if !server.origin_allowed(origin, host, scheme) {
return (StatusCode::FORBIDDEN, "origin not allowed").into_response();
}
}
next.run(req).await
}
async fn serve_mcp(
axum::extract::Extension(server): axum::extract::Extension<Arc<McpServer>>,
identity: Option<axum::extract::Extension<crate::security::ResolvedClientIdentity>>,
connect_info: Option<
axum::extract::Extension<axum::extract::ConnectInfo<std::net::SocketAddr>>,
>,
#[cfg(feature = "tls")] client_cert: Option<
axum::extract::Extension<std::sync::Arc<crate::tls::client_auth::ClientIdentity>>,
>,
envelope_admission: Option<axum::extract::Extension<crate::middleware::EnvelopeAdmission>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
let identity = identity.as_ref().map(|ext| &ext.0);
#[cfg(feature = "tls")]
let client_cert = client_cert.map(|ext| ext.0);
let origin = headers
.get(header::ORIGIN)
.and_then(|o| o.to_str().ok())
.map(str::to_owned);
let parsed: Value = match serde_json::from_slice(&body) {
Ok(v) => v,
Err(e) => {
if modern::header_is_modern(&headers) {
let rejection = modern::parse_error(&e.to_string());
let mut response = json_response(&rejection.body);
*response.status_mut() = rejection.status;
return response;
}
return json_response(&parse_error(&e.to_string()));
}
};
let modern = modern::wants_modern(&headers, &parsed);
if !modern && let Some(rejection) = reject_unsupported_protocol_version(&headers, &parsed) {
return rejection;
}
let ctx = ReplayContext {
headers: &headers,
identity,
peer: connect_info.map(|ext| (ext.0).0),
envelope_admission: envelope_admission.map(|ext| ext.0),
#[cfg(feature = "tls")]
client_cert,
modern,
};
let mut response = if modern {
serve_modern(&server, &ctx, parsed).await
} else {
serve_legacy(&server, &ctx, parsed).await
};
apply_cors_headers(&server.cors, origin.as_deref(), &mut response);
response
}
async fn serve_legacy(server: &Arc<McpServer>, ctx: &ReplayContext<'_>, parsed: Value) -> Response {
match parsed {
Value::Array(batch) if batch.is_empty() => {
json_response(&error(Value::Null, -32600, "Invalid Request: empty batch"))
}
Value::Array(batch)
if batch.iter().any(|msg| {
msg.get("method")
.and_then(Value::as_str)
.is_some_and(|m| m == "tools/call" || is_task_method(m))
}) =>
{
json_response(&error(
Value::Null,
-32600,
"Invalid Request: batched tools/call or tasks/* is not supported; \
send each one as a single JSON-RPC request",
))
}
Value::Array(batch) => {
let mut out = Vec::new();
for msg in batch {
if let Some(resp) = handle_message(server, &msg) {
out.push(resp);
}
}
if out.is_empty() {
StatusCode::ACCEPTED.into_response()
} else {
json_response(&Value::Array(out))
}
}
msg @ Value::Object(_) => {
if let Some((id, params)) = single_request(&msg, "tools/call") {
serve_tools_call(server, ctx, id, params).await
} else if let Some((method, id, params)) = single_task_request(&msg) {
json_response(&serve_task_method(server, ctx, &method, id, ¶ms))
} else {
handle_message(server, &msg).map_or_else(
|| StatusCode::ACCEPTED.into_response(),
|v| json_response(&v),
)
}
}
_ => json_response(&error(
Value::Null,
-32600,
"Invalid Request: expected a JSON object or array",
)),
}
}
async fn serve_modern(server: &Arc<McpServer>, ctx: &ReplayContext<'_>, parsed: Value) -> Response {
let reply = |status: StatusCode, mut message: Value| {
modern::stamp(&mut message, &server.server_info());
let mut response = json_response(&message);
*response.status_mut() = status;
response
};
let rejected = |rejection: modern::Rejection| reply(rejection.status, rejection.body);
if parsed.is_array() {
return rejected(modern::batch_rejection());
}
if modern::is_notification(&parsed) {
return StatusCode::ACCEPTED.into_response();
}
if let Err(rejection) = modern::validate(ctx.headers, &parsed, |m| server.serves(m)) {
return rejected(rejection);
}
let id = parsed.get("id").cloned().unwrap_or(Value::Null);
let method = parsed.get("method").and_then(Value::as_str).unwrap_or("");
let params = parsed.get("params").cloned().unwrap_or(Value::Null);
match method {
"server/discover" => reply(StatusCode::OK, success(id, discover_result(server))),
"tools/list" => {
let mut result = tools_list_result(server);
modern::add_cache_hints(&mut result);
reply(StatusCode::OK, success(id, result))
}
"tools/call" => serve_tools_call(server, ctx, id, params).await,
m if is_task_method(m) => {
let message = serve_task_method(server, ctx, m, id, ¶ms);
reply(modern::status_for(&message), message)
}
other => rejected(modern::method_not_found(id, other)),
}
}
fn reject_unsupported_protocol_version(headers: &HeaderMap, parsed: &Value) -> Option<Response> {
let is_initialize = parsed
.as_object()
.and_then(|o| o.get("method"))
.and_then(Value::as_str)
== Some("initialize");
if is_initialize {
return None;
}
let version = headers.get("mcp-protocol-version")?.to_str().unwrap_or("");
if SUPPORTED_PROTOCOL_VERSIONS.contains(&version) {
return None;
}
Some(
(
StatusCode::BAD_REQUEST,
format!("unsupported MCP-Protocol-Version: {version}"),
)
.into_response(),
)
}
fn handle_message(server: &McpServer, msg: &Value) -> Option<Value> {
let id = msg.get("id").cloned();
let id_ok = id
.as_ref()
.is_none_or(|v| v.is_string() || v.is_number() || v.is_null());
let is_valid = msg.is_object()
&& msg.get("jsonrpc").and_then(Value::as_str) == Some("2.0")
&& msg.get("method").and_then(Value::as_str).is_some()
&& id_ok;
if !is_valid {
let err_id = match &id {
Some(v) if v.is_string() || v.is_number() => v.clone(),
_ => Value::Null,
};
return Some(error(err_id, -32600, "Invalid Request"));
}
let id = id?;
let method = msg.get("method").and_then(Value::as_str).unwrap_or("");
let params = msg.get("params").cloned().unwrap_or(Value::Null);
let result = match method {
"initialize" => Ok(initialize_result(server, ¶ms)),
"ping" => Ok(json!({})),
"tools/list" => Ok(tools_list_result(server)),
"tools/call" => Err((
-32600,
"tools/call must be sent as a single JSON-RPC request".to_owned(),
)),
other => Err((-32601, format!("method not found: {other}"))),
};
Some(match result {
Ok(value) => success(id, value),
Err((code, message)) => error(id, code, &message),
})
}
impl McpServer {
fn capabilities(&self) -> Value {
let mut capabilities = json!({ "tools": { "listChanged": false } });
if self.has_task_tools {
capabilities["extensions"] = json!({ tasks::EXTENSION_ID: {} });
}
capabilities
}
fn serves(&self, method: &str) -> bool {
matches!(method, "server/discover" | "tools/list" | "tools/call")
|| (is_task_method(method) && self.has_task_tools)
}
fn server_info(&self) -> Value {
json!({ "name": self.server_name, "version": self.server_version })
}
}
fn discover_result(server: &McpServer) -> Value {
let mut result = json!({
"supportedVersions": modern::VERSIONS,
"capabilities": server.capabilities(),
});
if let Some(text) = &server.instructions {
result["instructions"] = json!(text);
}
modern::add_cache_hints(&mut result);
result
}
fn initialize_result(server: &McpServer, params: &Value) -> Value {
let protocol = match params.get("protocolVersion").and_then(Value::as_str) {
Some(requested) if SUPPORTED_PROTOCOL_VERSIONS.contains(&requested) => requested,
_ => DEFAULT_PROTOCOL_VERSION,
};
let mut result = json!({
"protocolVersion": protocol,
"capabilities": server.capabilities(),
"serverInfo": {
"name": server.server_name,
"version": server.server_version,
},
});
if let Some(text) = &server.instructions {
result["instructions"] = json!(text);
}
result
}
fn tools_list_result(server: &McpServer) -> Value {
let tools: Vec<Value> = match server.catalog {
McpCatalog::Progressive => server
.meta_tools
.iter()
.cloned()
.chain(
server
.tools
.iter()
.filter(|t| t.core)
.map(McpTool::descriptor),
)
.collect(),
McpCatalog::Full => server.tools.iter().map(McpTool::descriptor).collect(),
};
json!({ "tools": tools })
}
fn single_request(msg: &Value, method: &str) -> Option<(Value, Value)> {
let obj = msg.as_object()?;
if obj.get("jsonrpc").and_then(Value::as_str) != Some("2.0") {
return None;
}
if obj.get("method").and_then(Value::as_str) != Some(method) {
return None;
}
let id = obj.get("id")?;
if !(id.is_string() || id.is_number() || id.is_null()) {
return None;
}
let params = obj.get("params").cloned().unwrap_or(Value::Null);
Some((id.clone(), params))
}
const TASK_METHODS: &[&str] = &["tasks/get", "tasks/update", "tasks/cancel"];
fn is_task_method(method: &str) -> bool {
TASK_METHODS.contains(&method)
}
fn single_task_request(msg: &Value) -> Option<(String, Value, Value)> {
let method = msg.get("method").and_then(Value::as_str)?;
if !is_task_method(method) {
return None;
}
let (id, params) = single_request(msg, method)?;
Some((method.to_owned(), id, params))
}
impl McpServer {
fn task_owner(&self, ctx: &ReplayContext<'_>) -> tasks::Owner {
let headers = ctx.headers;
let header = |name: &str| headers.get(name).and_then(|v| v.to_str().ok());
let session = headers
.get_all(header::COOKIE)
.iter()
.filter_map(|v| v.to_str().ok())
.flat_map(|v| v.split(';'))
.filter_map(|pair| pair.trim().split_once('='))
.find(|(name, _)| *name == self.session_cookie)
.map(|(_, value)| value);
#[cfg(feature = "tls")]
let certificate = ctx.client_cert.as_deref().map(|c| c.fingerprint.as_str());
#[cfg(not(feature = "tls"))]
let certificate = None;
tasks::Owner::new(&[
header("authorization"),
session,
self.tenant_header.as_deref().and_then(header),
certificate,
])
}
fn new_task_id(&self) -> String {
let mut bytes = [0u8; 32];
self.state.entropy().fill_bytes(&mut bytes);
bytes.iter().fold(String::with_capacity(64), |mut out, b| {
let _ = write!(out, "{b:02x}");
out
})
}
}
fn serve_task_method(
server: &McpServer,
ctx: &ReplayContext<'_>,
method: &str,
id: Value,
params: &Value,
) -> Value {
if !server.has_task_tools {
return error(id, -32601, &format!("method not found: {method}"));
}
if !tasks::client_declares_tasks(params) {
return error_with_data(
id,
tasks::MISSING_CAPABILITY,
"Missing required client capability",
tasks::missing_capability_data(),
);
}
let Some(task_id) = params.get("taskId").and_then(Value::as_str) else {
return error(id, -32602, "`taskId` must be a string");
};
let owner = server.task_owner(ctx);
let now = server.state.clock().now();
let ack = || json!({ "resultType": "complete" });
let found = match method {
"tasks/get" => server.tasks.get(task_id, &owner, now),
"tasks/cancel" => server.tasks.cancel(task_id, &owner, now).then(ack),
_ if !params.get("inputResponses").is_some_and(Value::is_object) => {
return error(id, -32602, "`inputResponses` must be a JSON object");
}
_ => server.tasks.contains(task_id, &owner, now).then(ack),
};
found.map_or_else(
|| {
error(
id.clone(),
-32602,
"Failed to retrieve task: Task not found",
)
},
|result| success(id.clone(), result),
)
}
const MAX_AUDITED_EFFECTS: usize = 16;
const MAX_ARGUMENT_NAMES_BYTES: usize = 512;
const ANONYMOUS_AGENT: &str = "agent:anonymous";
async fn admit_call(server: &McpServer, tool: &McpTool, audit: &AgentAudit) -> Result<(), String> {
if let Err(refusal) = check_rate_cap(server, tool) {
audit.record_refusal(&server.state, &refusal).await;
return Err(refusal);
}
audit
.record_attempt(&server.state)
.await
.map_err(str::to_owned)
}
fn check_rate_cap(server: &McpServer, tool: &McpTool) -> Result<(), String> {
let Some(grant) = tool.agent_authority.map(|a| a.grant) else {
return Ok(());
};
let Some(rate) = grant.rate else {
return Ok(());
};
let Some(cap) = crate::agent_authority::rate::parse_rate(rate) else {
return Err(format!(
"rate cap `{rate}` for grant {} cannot be read; the call is refused",
grant.name
));
};
server
.rate_limiter
.check(grant.name, grant.location, cap, server.state.clock().now())
.map_err(|wait| {
format!(
"rate cap {rate} for grant {} exceeded; retry in {wait}s",
grant.name
)
})
}
const MAX_ACTOR_LEN: usize = 128;
fn clean_actor(actor: &str) -> Option<String> {
let clean: String = actor
.chars()
.filter(|c| {
c.is_alphanumeric()
|| matches!(
c,
'-' | '_' | '.' | ':' | '@' | '+' | '/' | '=' | '#' | '~' | ' '
)
})
.take(MAX_ACTOR_LEN)
.collect();
(!clean.trim().is_empty()).then_some(clean)
}
const UNAUDITABLE_REFUSAL: &str =
"audit attempt record could not be written; refusing a non-reversible action";
const AUDIT_WRITE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2);
#[derive(Clone)]
struct AgentAudit {
tool: String,
path_template: String,
correlation_id: String,
reversibility: Option<crate::agent_authority::Reversibility>,
grant: Option<&'static str>,
actor: String,
ip: Option<std::net::IpAddr>,
shared: std::collections::BTreeMap<String, String>,
}
impl AgentAudit {
fn new(
server: &McpServer,
tool: &McpTool,
arguments: &Value,
ctx: &ReplayContext<'_>,
via_call_tool: bool,
) -> Self {
let authority = tool.agent_authority;
let grant = authority.map(|a| a.grant.name);
let reversibility = authority.map(|a| a.grant.reversibility);
let correlation_id = server.state.entropy().uuid_v4().to_string();
let mut shared = std::collections::BTreeMap::new();
shared.insert("correlation_id".to_owned(), correlation_id.clone());
shared.insert("transport".to_owned(), "mcp".to_owned());
shared.insert("tool".to_owned(), tool.name.clone());
if let Some(grant) = grant {
shared.insert("grant".to_owned(), grant.to_owned());
}
shared.insert(
"reversibility".to_owned(),
reversibility.map_or_else(|| "unknown".to_owned(), |r| r.as_str().to_owned()),
);
if let Some(authority) = authority {
shared.insert("effects".to_owned(), render_effects(authority.effects));
}
shared.insert("argument_names".to_owned(), argument_names(tool, arguments));
if via_call_tool {
shared.insert("via".to_owned(), catalog::CALL_TOOL.to_owned());
}
Self {
tool: tool.name.clone(),
path_template: tool.path_template.clone(),
correlation_id,
reversibility,
grant,
actor: crate::current::Current::scoped_actor()
.and_then(|actor| clean_actor(&actor))
.unwrap_or_else(|| ANONYMOUS_AGENT.to_owned()),
ip: ctx.peer.map(|peer| peer.ip()),
shared,
}
}
fn invocation(&self) -> crate::agent_authority::AgentInvocation {
crate::agent_authority::AgentInvocation {
correlation_id: self.correlation_id.clone(),
tool: self.tool.clone(),
grant: self.grant,
reversibility: self.reversibility,
}
}
fn refuse_when_unauditable(&self) -> bool {
self.reversibility != Some(crate::agent_authority::Reversibility::Reversible)
}
fn adopt_actor(&mut self, actor: Option<&str>) {
if self.actor != ANONYMOUS_AGENT {
return;
}
if let Some(clean) = actor.and_then(clean_actor) {
self.actor = clean;
}
}
fn event(
&self,
action: String,
status: crate::audit::AuditStatus,
phase: &str,
) -> crate::audit::AuditEvent {
let mut event = crate::audit::AuditEvent::new(
self.actor.clone(),
action,
self.path_template.clone(),
self.ip,
status,
);
event.metadata = self.shared.clone();
event.metadata.insert("phase".to_owned(), phase.to_owned());
event
}
async fn write_bounded(
state: &crate::state::AppState,
event: crate::audit::AuditEvent,
) -> Result<(), String> {
match tokio::time::timeout(
AUDIT_WRITE_TIMEOUT,
crate::audit::write_from_state(state, event),
)
.await
{
Ok(Ok(())) => Ok(()),
Ok(Err(error)) => Err(error.to_string()),
Err(_elapsed) => Err(format!(
"audit sink did not respond within {AUDIT_WRITE_TIMEOUT:?}"
)),
}
}
fn reversibility_field(&self) -> Option<&str> {
self.shared.get("reversibility").map(String::as_str)
}
async fn record_attempt(&self, state: &crate::state::AppState) -> Result<(), &'static str> {
let event = self.event(
format!("agent.tool.{}.attempt", self.tool),
crate::audit::AuditStatus::Success,
"attempt",
);
let Err(error) = Self::write_bounded(state, event).await else {
tracing::info!(
target: "autumn.agent",
tool = %self.tool,
correlation_id = %self.correlation_id,
transport = "mcp",
phase = "attempt",
grant = self.grant,
reversibility = self.reversibility_field(),
effects = self.shared.get("effects").map(String::as_str),
argument_names = self.shared.get("argument_names").map(String::as_str),
actor = %self.actor,
target = %self.path_template,
"agent tool call attempted"
);
return Ok(());
};
if self.refuse_when_unauditable() {
tracing::error!(
target: "autumn.agent",
tool = %self.tool,
correlation_id = %self.correlation_id,
transport = "mcp",
phase = "refused",
grant = self.grant,
reversibility = self.reversibility_field(),
actor = %self.actor,
target = %self.path_template,
error = %error,
"refusing an agent tool call: its attempt record could not be written \
and the action is not reversible"
);
self.record_refusal(state, &error).await;
return Err(UNAUDITABLE_REFUSAL);
}
tracing::warn!(
target: "autumn.agent",
tool = %self.tool,
correlation_id = %self.correlation_id,
phase = "attempt",
error = %error,
"agent tool attempt record could not be written; the action is reversible, \
so the call proceeds"
);
Ok(())
}
async fn record_refusal(&self, state: &crate::state::AppState, error: &str) {
let mut event = self.event(
format!("agent.tool.{}.refused", self.tool),
crate::audit::AuditStatus::Failure,
"refused",
);
event
.metadata
.insert("refused_reason".to_owned(), error.to_owned());
let _ = Self::write_bounded(state, event).await;
}
async fn record_outcome(
&self,
state: &crate::state::AppState,
status: Option<StatusCode>,
request_id: Option<&str>,
disposition: Disposition,
) {
let succeeded = status.is_some_and(|s| s.is_success()) && disposition.is_success();
let audit_status = if succeeded {
crate::audit::AuditStatus::Success
} else {
crate::audit::AuditStatus::Failure
};
let mut event = self.event(format!("agent.tool.{}", self.tool), audit_status, "outcome");
if let Some(status) = status {
event
.metadata
.insert("http_status".to_owned(), status.as_u16().to_string());
}
if let Some(stream) = disposition.stream_state() {
event
.metadata
.insert("stream_state".to_owned(), stream.to_owned());
}
if let Some(result) = disposition.result() {
event
.metadata
.insert("result".to_owned(), result.to_owned());
}
if let Some(request_id) = request_id {
event
.metadata
.insert("request_id".to_owned(), request_id.to_owned());
}
tracing::info!(
target: "autumn.agent",
tool = %self.tool,
correlation_id = %self.correlation_id,
transport = "mcp",
phase = "outcome",
grant = self.grant,
reversibility = self.reversibility_field(),
effects = self.shared.get("effects").map(String::as_str),
argument_names = self.shared.get("argument_names").map(String::as_str),
actor = %self.actor,
target = %self.path_template,
http_status = status.map(|s| s.as_u16()),
stream_state = disposition.stream_state(),
result = disposition.result(),
request_id = request_id,
"agent tool call"
);
if let Err(error) = Self::write_bounded(state, event).await {
tracing::warn!(
target: "autumn.agent",
tool = %self.tool,
correlation_id = %self.correlation_id,
phase = "outcome",
error = %error,
"agent tool outcome record could not be written"
);
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BufferedFailure {
Overflow,
BodyError,
}
impl BufferedFailure {
fn classify(error: &axum::Error) -> Self {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error);
while let Some(err) = source {
if err.is::<http_body_util::LengthLimitError>() {
return Self::Overflow;
}
source = err.source();
}
Self::BodyError
}
const fn as_str(self) -> &'static str {
match self {
Self::Overflow => "body_overflow",
Self::BodyError => "body_error",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Disposition {
Settled,
Stream(StreamState),
Buffered(BufferedFailure),
Task(TaskEnd),
}
impl Disposition {
const fn is_success(self) -> bool {
match self {
Self::Settled => true,
Self::Stream(state) => state.is_success(),
Self::Buffered(_) | Self::Task(_) => false,
}
}
const fn stream_state(self) -> Option<&'static str> {
match self {
Self::Stream(state) => Some(state.as_str()),
_ => None,
}
}
const fn result(self) -> Option<&'static str> {
match self {
Self::Buffered(failure) => Some(failure.as_str()),
Self::Task(end) => Some(end.as_str()),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TaskEnd {
Aborted,
Panicked,
}
impl TaskEnd {
const fn as_str(self) -> &'static str {
match self {
Self::Aborted => "task_aborted",
Self::Panicked => "task_panicked",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StreamState {
Completed,
Aborted,
Errored,
}
impl StreamState {
const fn as_str(self) -> &'static str {
match self {
Self::Completed => "completed",
Self::Aborted => "aborted",
Self::Errored => "errored",
}
}
const fn is_success(self) -> bool {
matches!(self, Self::Completed)
}
}
struct OutcomeGuard {
audit: AgentAudit,
state: crate::state::AppState,
status: Option<StatusCode>,
request_id: Option<String>,
abandoned: Disposition,
recorded: bool,
}
impl OutcomeGuard {
const fn new(audit: AgentAudit, state: crate::state::AppState, abandoned: Disposition) -> Self {
Self {
audit,
state,
status: None,
request_id: None,
abandoned,
recorded: false,
}
}
async fn record(&mut self, disposition: Disposition) {
if self.recorded {
return;
}
self.recorded = true;
self.audit
.record_outcome(
&self.state,
self.status,
self.request_id.as_deref(),
disposition,
)
.await;
}
}
impl Drop for OutcomeGuard {
fn drop(&mut self) {
if self.recorded {
return;
}
self.recorded = true;
let disposition = match self.abandoned {
Disposition::Task(_) if std::thread::panicking() => {
Disposition::Task(TaskEnd::Panicked)
}
other => other,
};
let audit = self.audit.clone();
let state = self.state.clone();
let status = self.status;
let request_id = self.request_id.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
audit
.record_outcome(&state, status, request_id.as_deref(), disposition)
.await;
});
}
}
}
fn render_effects(effects: &[crate::agent_authority::Effect]) -> String {
let mut out = String::new();
for effect in effects.iter().take(MAX_AUDITED_EFFECTS) {
if !out.is_empty() {
out.push(',');
}
out.push_str(effect.kind.as_str());
out.push(':');
out.push_str(effect.subject);
}
if effects.len() > MAX_AUDITED_EFFECTS {
out.push_str(",…");
}
out
}
fn argument_names(tool: &McpTool, arguments: &Value) -> String {
let Some(object) = arguments.as_object() else {
return String::new();
};
let is_known = |name: &str| {
name == "body"
|| name == "query"
|| tool
.path_params
.iter()
.any(|param| param.strip_prefix('*').unwrap_or(param) == name)
};
let mut known: Vec<&str> = Vec::new();
let mut unknown: usize = 0;
for name in object.keys() {
if is_known(name) {
known.push(name);
} else {
unknown += 1;
}
}
known.sort_unstable();
let mut out = String::new();
for name in known {
if out.len() + name.len() + 1 > MAX_ARGUMENT_NAMES_BYTES {
out.push_str(",…");
break;
}
if !out.is_empty() {
out.push(',');
}
out.push_str(name);
}
if unknown > 0 {
if !out.is_empty() {
out.push(',');
}
let _ = write!(out, "+{unknown} unknown");
}
out
}
fn apply_replay_extensions(
request: &mut axum::extract::Request,
server: &McpServer,
ctx: &ReplayContext<'_>,
audit: &AgentAudit,
detached: bool,
) {
let extensions = request.extensions_mut();
extensions.insert(audit.invocation());
if let Some(identity) = ctx.identity {
extensions.insert(identity.clone());
}
if let Some(peer) = ctx.peer {
extensions.insert(axum::extract::ConnectInfo(peer));
}
#[cfg(feature = "tls")]
if let Some(client_cert) = &ctx.client_cert {
extensions.insert(std::sync::Arc::clone(client_cert));
}
if server.envelope_rate_limited {
extensions.insert(crate::security::RateLimitEnvelopeCounted);
}
if server.envelope_load_shed && !detached {
extensions.insert(crate::middleware::LoadShedExempt);
if let Some(admission) = &ctx.envelope_admission {
extensions.insert(admission.clone());
}
}
}
struct Dispatched {
status: StatusCode,
request_id: Option<String>,
is_event_stream: bool,
replayed: Vec<(HeaderName, HeaderValue)>,
actor: Option<String>,
}
impl Dispatched {
fn inspect(response: &Response) -> Self {
let headers = response.headers();
Self {
actor: response
.extensions()
.get::<crate::agent_authority::AgentActor>()
.map(|actor| actor.0.clone()),
status: response.status(),
request_id: headers
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.map(str::to_owned),
is_event_stream: headers
.get(header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.is_some_and(|c| {
c.trim_start()
.to_ascii_lowercase()
.starts_with("text/event-stream")
}),
replayed: REPLAYED_RESPONSE_HEADERS
.iter()
.flat_map(|&name| {
headers
.get_all(name)
.iter()
.map(move |value| (HeaderName::from_static(name), value.clone()))
})
.collect(),
}
}
}
fn apply_replayed(headers: &mut HeaderMap, replayed: Vec<(HeaderName, HeaderValue)>) {
let mut paid = false;
for (name, value) in replayed {
paid |= name.as_str().starts_with("payment-");
headers.append(name, value);
}
if paid {
headers.insert(
header::CACHE_CONTROL,
HeaderValue::from_static("private, no-store"),
);
crate::aeo::commerce::expose_payment_headers(headers);
}
}
async fn serve_tools_call(
server: &Arc<McpServer>,
ctx: &ReplayContext<'_>,
id: Value,
params: Value,
) -> Response {
let server_info = ctx.modern.then(|| server.server_info());
let reply = |mut message: Value| {
if let Some(info) = &server_info {
modern::stamp(&mut message, info);
}
json_response(&message)
};
let (params, arguments, via_call_tool) = match prepare_call(server, &id, params, reply) {
PreparedCall::Run(params, arguments, via_call_tool) => (params, arguments, via_call_tool),
PreparedCall::Done(response) => return response,
};
let name = params.get("name").and_then(Value::as_str).unwrap_or("");
let Some(&idx) = server.by_name.get(name) else {
return json_response(&error(id, -32602, &format!("unknown tool: {name}")));
};
let tool = &server.tools[idx];
let mut request = match build_request(
tool,
ctx.headers,
&arguments,
&server.csrf_header,
server.tenant_header.as_deref(),
) {
Ok(req) => req,
Err(message) => return json_response(&error(id, -32602, &message)),
};
let mut audit = AgentAudit::new(server, tool, &arguments, ctx, via_call_tool);
if let Err(refusal) = admit_call(server, tool, &audit).await {
return reply(success(id, tool_error(&refusal)));
}
if tool.task
&& tasks::client_declares_tasks(¶ms)
&& !crate::aeo::commerce::prices_route(&server.state, &tool.method, &tool.path_template)
&& let Some((task_id, created)) = reserve_task(server, ctx, &tool.name)
{
apply_replay_extensions(&mut request, server, ctx, &audit, true);
spawn_task(server, idx, &task_id, request, audit);
return reply(success(id, created));
}
apply_replay_extensions(&mut request, server, ctx, &audit, false);
let response = match server.dispatch.clone().oneshot(request).await {
Ok(resp) => resp,
Err(e) => {
audit
.record_outcome(
&server.state,
Some(StatusCode::INTERNAL_SERVER_ERROR),
None,
Disposition::Settled,
)
.await;
return reply(success(id, tool_error(&format!("dispatch failed: {e}"))));
}
};
let Dispatched {
status,
request_id,
is_event_stream,
replayed,
actor,
} = Dispatched::inspect(&response);
audit.adopt_actor(actor.as_deref());
let client_accepts_sse = ctx
.headers
.get(header::ACCEPT)
.and_then(|v| v.to_str().ok())
.is_some_and(accept_includes_event_stream);
if tool.streams && status.is_success() && is_event_stream && client_accepts_sse {
let mut stream_audit = OutcomeGuard::new(
audit,
server.state.clone(),
Disposition::Stream(StreamState::Aborted),
);
stream_audit.status = Some(status);
stream_audit.request_id = request_id;
return stream_tool_result(id, ¶ms, response, replayed, stream_audit, server_info);
}
let (resp, failure) = buffered_tool_response(
tool,
id,
status,
is_event_stream,
replayed,
response,
server_info.as_ref(),
)
.await;
let disposition = failure.map_or(Disposition::Settled, Disposition::Buffered);
audit
.record_outcome(
&server.state,
Some(status),
request_id.as_deref(),
disposition,
)
.await;
resp
}
enum PreparedCall {
Run(Value, Value, bool),
Done(Response),
}
fn prepare_call(
server: &McpServer,
id: &Value,
params: Value,
reply: impl Fn(Value) -> Response,
) -> PreparedCall {
let invalid =
|message: &str| PreparedCall::Done(json_response(&error(id.clone(), -32602, message)));
let (params, via_call_tool) = match params.get("name").and_then(Value::as_str) {
Some(catalog::SEARCH_TOOL) if server.catalog == McpCatalog::Progressive => {
return match search_tools(server, ¶ms) {
Ok(result) => PreparedCall::Done(reply(success(id.clone(), result))),
Err(message) => invalid(&message),
};
}
Some(catalog::CALL_TOOL) if server.catalog == McpCatalog::Progressive => {
match catalog::unwrap_call(¶ms) {
Ok(inner) => (inner, true),
Err(message) => return invalid(&message),
}
}
_ => (params, false),
};
match call_arguments(¶ms) {
Ok(arguments) => PreparedCall::Run(params, arguments, via_call_tool),
Err(message) => invalid(message),
}
}
fn call_arguments(params: &Value) -> Result<Value, &'static str> {
match params.get("arguments") {
None => Ok(json!({})),
Some(value) if value.is_object() => Ok(value.clone()),
Some(_) => Err("`arguments` must be a JSON object"),
}
}
fn search_tools(server: &McpServer, params: &Value) -> Result<Value, String> {
let arguments = call_arguments(params)?;
let args = catalog::parse_search(&arguments)?;
let found = catalog::search(&server.tools, &server.search_index, &args);
let mut result = tool_ok(&found.to_string());
result["structuredContent"] = found;
Ok(result)
}
fn reserve_task(
server: &McpServer,
ctx: &ReplayContext<'_>,
tool: &str,
) -> Option<(String, Value)> {
let task_id = server.new_task_id();
let owner = server.task_owner(ctx);
let quota = ctx
.identity
.and_then(|id| id.addr)
.or_else(|| ctx.peer.map(|peer| peer.ip()))
.map_or_else(|| "unknown".to_owned(), |ip| ip.to_string());
let now = server.state.clock().now();
let Some(created) = server.tasks.create(&task_id, owner, "a, now) else {
tracing::warn!(
tool = %tool,
"MCP task store is full; the tool call runs as a blocking call"
);
return None;
};
Some((task_id, created))
}
fn spawn_task(
server: &Arc<McpServer>,
idx: usize,
task_id: &str,
request: axum::extract::Request,
audit: AgentAudit,
) {
let outcome = OutcomeGuard::new(
audit,
server.state.clone(),
Disposition::Task(TaskEnd::Aborted),
);
let ttl = server.tasks.ttl();
let job_server = Arc::clone(server);
let job_id = task_id.to_owned();
let job = tokio::spawn(async move {
let server = job_server;
let run = std::panic::AssertUnwindSafe(run_task(&server, idx, request, outcome));
let Ok(ran) = tokio::time::timeout(ttl, run.catch_unwind()).await else {
let error = json!({ "code": -32603, "message": "Internal error" });
let state = tasks::TaskState::Failed(error, "The task TTL ended.".to_owned());
server
.tasks
.finish(&job_id, state, server.state.clock().now());
return;
};
let Ok((mut result, mut outcome, disposition, permit)) = ran else {
let error = json!({ "code": -32603, "message": "Internal error" });
let state = tasks::TaskState::Failed(error, "The tool handler panicked.".to_owned());
server
.tasks
.finish(&job_id, state, server.state.clock().now());
return;
};
if let Some(object) = result.as_object_mut() {
object.insert("resultType".to_owned(), json!("complete"));
}
let state = tasks::TaskState::Completed(Arc::new(result));
server
.tasks
.finish(&job_id, state, server.state.clock().now());
drop(permit);
outcome.record(disposition).await;
});
server.tasks.set_abort(task_id, &job.abort_handle());
let sweep_server = Arc::clone(server);
tokio::spawn(async move {
tokio::time::sleep(ttl).await;
sweep_server.tasks.sweep(sweep_server.state.clock().now());
});
}
async fn run_task(
server: &McpServer,
idx: usize,
request: axum::extract::Request,
mut outcome: OutcomeGuard,
) -> (
Value,
OutcomeGuard,
Disposition,
Option<tokio::sync::SemaphorePermit<'_>>,
) {
let tool = &server.tools[idx];
let response = match server.dispatch.clone().oneshot(request).await {
Ok(resp) => resp,
Err(e) => {
outcome.status = Some(StatusCode::INTERNAL_SERVER_ERROR);
let result = tool_error(&format!("dispatch failed: {e}"));
return (result, outcome, Disposition::Settled, None);
}
};
let Dispatched {
status,
request_id,
is_event_stream,
actor,
..
} = Dispatched::inspect(&response);
outcome.audit.adopt_actor(actor.as_deref());
outcome.status = Some(status);
outcome.request_id = request_id;
let permit = server.tasks.buffer_permit().await;
let (result, failure) =
read_tool_result(tool, status, is_event_stream, response.into_body()).await;
let disposition = failure.map_or(Disposition::Settled, Disposition::Buffered);
(result, outcome, disposition, permit)
}
async fn buffered_tool_response(
tool: &McpTool,
id: Value,
status: StatusCode,
is_event_stream: bool,
replayed: Vec<(HeaderName, HeaderValue)>,
response: Response,
server_info: Option<&Value>,
) -> (Response, Option<BufferedFailure>) {
let mut server_timings: Vec<HeaderValue> = Vec::new();
for value in response.headers().get_all(&SERVER_TIMING) {
server_timings.push(value.clone());
}
let (value, failure) =
read_tool_result(tool, status, is_event_stream, response.into_body()).await;
let mut message = success(id, value);
if let Some(info) = server_info {
modern::stamp(&mut message, info);
}
let mut resp = json_response(&message);
apply_replayed(resp.headers_mut(), replayed);
if failure.is_some() {
return (resp, failure);
}
for timing in server_timings {
if let Ok(value) = timing.to_str()
&& let Some(kept) = strip_total_metric(value)
&& let Ok(hv) = HeaderValue::from_str(&kept)
{
resp.headers_mut().append(SERVER_TIMING.clone(), hv);
}
}
(resp, None)
}
async fn read_tool_result(
tool: &McpTool,
status: StatusCode,
is_event_stream: bool,
body: Body,
) -> (Value, Option<BufferedFailure>) {
let bytes = match axum::body::to_bytes(body, MAX_TOOL_RESPONSE_BYTES).await {
Ok(bytes) => bytes,
Err(error) => {
let failure = BufferedFailure::classify(&error);
let message = match failure {
BufferedFailure::Overflow => format!(
"handler response exceeded the {MAX_TOOL_RESPONSE_BYTES}-byte MCP tool-result limit"
),
BufferedFailure::BodyError => {
format!("handler response body failed mid-read: {error}")
}
};
return (tool_error(&message), Some(failure));
}
};
let text = if is_event_stream {
collapse_sse_body(&bytes)
} else {
String::from_utf8_lossy(&bytes).into_owned()
};
(buffered_tool_result(tool, status, &text), None)
}
fn strip_total_metric(value: &str) -> Option<String> {
let kept: Vec<&str> = value
.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.filter(|entry| {
let name = entry.split([';', '=']).next().unwrap_or("").trim();
!name.eq_ignore_ascii_case("total")
})
.collect();
if kept.is_empty() {
None
} else {
Some(kept.join(", "))
}
}
fn buffered_tool_result(tool: &McpTool, status: StatusCode, text: &str) -> Value {
if !status.is_success() {
return tool_error(&format!(
"handler returned HTTP {}: {text}",
status.as_u16()
));
}
if tool.empty_body {
if !text.is_empty() {
tracing::warn!(
tool = %tool.name,
body_len = text.len(),
"empty-body-contract tool returned a non-empty body; \
discarding it (the route's declared 204/205 status does \
not match what its handler actually returns)"
);
}
return tool_ok("");
}
tool_ok(text)
}
fn accept_includes_event_stream(accept: &str) -> bool {
accept.split(',').any(|part| {
let media = part
.split(';')
.next()
.unwrap_or("")
.trim()
.to_ascii_lowercase();
media == "text/event-stream" || media == "text/*"
})
}
fn progress_notification(token: &Value, progress: f64, message: &str) -> Value {
if let Ok(Value::Object(map)) = serde_json::from_str::<Value>(message)
&& map.get("progress").is_some_and(Value::is_number)
{
let mut params = serde_json::Map::new();
params.insert("progressToken".into(), token.clone());
for key in ["progress", "total", "message"] {
if let Some(v) = map.get(key) {
params.insert((*key).to_owned(), v.clone());
}
}
return json!({
"jsonrpc": "2.0",
"method": "notifications/progress",
"params": Value::Object(params),
});
}
json!({
"jsonrpc": "2.0",
"method": "notifications/progress",
"params": { "progressToken": token, "progress": progress, "message": message },
})
}
enum ProjectionPhase {
Streaming,
Final,
Done,
}
struct StreamProjection {
body: Pin<Box<dyn Stream<Item = Result<Bytes, axum::Error>> + Send>>,
parser: SseWireParser,
ready: VecDeque<Event>,
progress_token: Option<Value>,
id: Value,
progress: f64,
progress_parts: Vec<String>,
result_parts: Vec<String>,
phase: ProjectionPhase,
audit: OutcomeGuard,
stamp: Option<Value>,
}
fn stream_tool_result(
id: Value,
params: &Value,
response: Response,
replayed: Vec<(HeaderName, HeaderValue)>,
audit: OutcomeGuard,
stamp: Option<Value>,
) -> Response {
let progress_token = params
.get("_meta")
.and_then(|m| m.get("progressToken"))
.cloned();
let state = StreamProjection {
body: Box::pin(response.into_body().into_data_stream()),
parser: SseWireParser::new(),
ready: VecDeque::new(),
progress_token,
id,
progress: 0.0,
progress_parts: Vec::new(),
result_parts: Vec::new(),
phase: ProjectionPhase::Streaming,
audit,
stamp,
};
let stream = futures::stream::unfold(state, project_next);
let mut resp = Sse::new(stream)
.keep_alive(crate::sse::keep_alive())
.into_response();
resp.headers_mut()
.insert("x-accel-buffering", HeaderValue::from_static("no"));
apply_replayed(resp.headers_mut(), replayed);
resp
}
async fn project_next(
mut st: StreamProjection,
) -> Option<(Result<Event, Infallible>, StreamProjection)> {
loop {
if let Some(event) = st.ready.pop_front() {
return Some((Ok(event), st));
}
match st.phase {
ProjectionPhase::Done => return None,
ProjectionPhase::Final => {
let content = if st.result_parts.is_empty() {
st.progress_parts.join("\n")
} else {
st.result_parts.concat()
};
let mut value = success(st.id.clone(), tool_ok(&content));
if let Some(server_info) = &st.stamp {
modern::stamp(&mut value, server_info);
}
st.phase = ProjectionPhase::Done;
return Some((Ok(Event::default().data(value.to_string())), st));
}
ProjectionPhase::Streaming => match st.body.next().await {
Some(Ok(bytes)) => {
let events = st.parser.push(&bytes);
enqueue_projected(&mut st, events);
}
terminal => {
let outcome = if terminal.is_some() {
StreamState::Errored
} else {
StreamState::Completed
};
let trailing = st.parser.finish();
enqueue_projected(&mut st, trailing);
st.phase = ProjectionPhase::Final;
st.audit.record(Disposition::Stream(outcome)).await;
}
},
}
}
}
fn enqueue_projected(st: &mut StreamProjection, events: Vec<ParsedSseEvent>) {
for ev in events {
if ev.event.as_deref() == Some("result") {
st.result_parts.push(ev.data);
continue;
}
st.progress_parts.push(ev.data.clone());
if let Some(token) = &st.progress_token {
st.progress += 1.0;
let note = progress_notification(token, st.progress, &ev.data);
st.ready.push_back(Event::default().data(note.to_string()));
}
}
}
struct ParsedSseEvent {
event: Option<String>,
data: String,
}
struct SseWireParser {
buffer: String,
event_type: Option<String>,
data_lines: Vec<String>,
has_fields: bool,
}
impl SseWireParser {
const fn new() -> Self {
Self {
buffer: String::new(),
event_type: None,
data_lines: Vec::new(),
has_fields: false,
}
}
fn push(&mut self, bytes: &[u8]) -> Vec<ParsedSseEvent> {
self.buffer.push_str(&String::from_utf8_lossy(bytes));
let mut out = Vec::new();
while let Some(pos) = self.buffer.find('\n') {
let line: String = self.buffer.drain(..=pos).collect();
if let Some(event) = self.process_line(line.trim_end_matches(['\n', '\r'])) {
out.push(event);
}
}
out
}
fn finish(&mut self) -> Vec<ParsedSseEvent> {
let mut out = Vec::new();
if !self.buffer.is_empty() {
let line = std::mem::take(&mut self.buffer);
if let Some(event) = self.process_line(line.trim_end_matches(['\n', '\r'])) {
out.push(event);
}
}
if let Some(event) = self.dispatch() {
out.push(event);
}
out
}
fn process_line(&mut self, line: &str) -> Option<ParsedSseEvent> {
if line.is_empty() {
return self.dispatch();
}
if line.starts_with(':') {
return None;
}
let (field, value) = match line.split_once(':') {
Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
None => (line, ""),
};
match field {
"event" => {
self.event_type = Some(value.to_owned());
self.has_fields = true;
}
"data" => {
self.data_lines.push(value.to_owned());
self.has_fields = true;
}
_ => {}
}
None
}
fn dispatch(&mut self) -> Option<ParsedSseEvent> {
if !self.has_fields {
return None;
}
let event = ParsedSseEvent {
event: self.event_type.take(),
data: self.data_lines.join("\n"),
};
self.data_lines.clear();
self.has_fields = false;
Some(event)
}
}
fn collapse_sse_body(bytes: &[u8]) -> String {
let mut parser = SseWireParser::new();
let mut events = parser.push(bytes);
events.extend(parser.finish());
let (results, progress): (Vec<_>, Vec<_>) = events
.into_iter()
.partition(|e| e.event.as_deref() == Some("result"));
if results.is_empty() {
progress
.into_iter()
.map(|e| e.data)
.collect::<Vec<_>>()
.join("\n")
} else {
results.into_iter().map(|e| e.data).collect()
}
}
fn build_request(
tool: &McpTool,
headers: &HeaderMap,
arguments: &Value,
csrf_header: &str,
tenant_header: Option<&str>,
) -> Result<axum::http::Request<Body>, String> {
let mut path = tool.path_template.clone();
for param in &tool.path_params {
let is_catch_all = param.starts_with('*');
let arg_key = param.strip_prefix('*').unwrap_or(param);
let raw = arguments
.get(arg_key)
.ok_or_else(|| format!("missing required path parameter `{arg_key}`"))?;
let value = match raw {
Value::String(s) => s.clone(),
Value::Number(_) | Value::Bool(_) => raw.to_string(),
_ => return Err(format!("path parameter `{arg_key}` must be a string")),
};
let encoded = if is_catch_all {
crate::paths::encode_catch_all_param(&value)
} else {
crate::paths::encode_path_segment(&value)
};
path = replace_path_param(&path, param, &encoded);
}
if tool.has_query
&& let Some(query) = arguments.get("query")
{
let Value::Object(map) = query else {
return Err("`query` must be a JSON object".to_owned());
};
let mut pairs = QueryPairs::new();
for (key, value) in map {
if !is_expressible_field_name(key) {
return Err(inexpressible_field_name("query"));
}
encode_query_arg(key, value, 1, &mut pairs)?;
}
if !pairs.pairs.is_empty() {
let qs = serde_urlencoded::to_string(&pairs.pairs)
.map_err(|e| format!("invalid query arguments: {e}"))?;
path = format!("{path}?{qs}");
}
}
let mut builder = axum::http::Request::builder()
.method(tool.method.as_str())
.uri(&path);
for name in FORWARDED_HEADERS {
for value in headers.get_all(*name) {
builder = builder.header(*name, value);
}
}
if let Some(value) = headers.get(csrf_header) {
builder = builder.header(csrf_header, value);
}
if let Some(name) = tenant_header
&& let Some(value) = headers.get(name)
{
builder = builder.header(name, value);
}
let body = if tool.has_body {
let payload = arguments
.get("body")
.ok_or_else(|| "missing required `body` argument".to_owned())?;
builder = builder.header(header::CONTENT_TYPE, "application/json");
Body::from(serde_json::to_vec(payload).unwrap_or_default())
} else {
Body::empty()
};
builder
.body(body)
.map_err(|e| format!("invalid request: {e}"))
}
fn query_scalar(value: &Value) -> String {
match value {
Value::String(s) => s.clone(),
other => other.to_string(),
}
}
const MAX_QUERY_PAIRS: usize = 1024;
const MAX_QUERY_BYTES: usize = 8 * 1024;
const MAX_QUERY_NODES: usize = 4 * MAX_QUERY_PAIRS;
struct QueryPairs {
pairs: Vec<(String, String)>,
bytes: usize,
nodes: usize,
}
impl QueryPairs {
const fn new() -> Self {
Self {
pairs: Vec::new(),
bytes: 0,
nodes: 0,
}
}
fn visit(&mut self) -> Result<(), String> {
self.nodes += 1;
if self.nodes > MAX_QUERY_NODES {
return Err(format!(
"query arguments traverse more than {MAX_QUERY_NODES} values; \
move the field to a JSON body"
));
}
Ok(())
}
fn push(&mut self, key: &str, value: String) -> Result<(), String> {
self.bytes = self
.bytes
.saturating_add(key.len())
.saturating_add(value.len());
if self.pairs.len() >= MAX_QUERY_PAIRS || self.bytes > MAX_QUERY_BYTES {
return Err(format!(
"query arguments expand past the dispatch limit of {MAX_QUERY_PAIRS} \
parameters / {MAX_QUERY_BYTES} bytes"
));
}
self.pairs.push((key.to_owned(), value));
Ok(())
}
}
fn is_expressible_field_name(name: &str) -> bool {
!name.is_empty() && !name.contains('[') && !name.contains(']')
}
fn inexpressible_field_name(owner: &str) -> String {
format!(
"query argument `{owner}` has a field name that a query string cannot carry: a name \
may not be empty or contain `[` or `]`, which are structure in the query encoding"
)
}
fn encode_query_arg(
key: &str,
value: &Value,
depth: usize,
out: &mut QueryPairs,
) -> Result<(), String> {
if depth > crate::query_string::MAX_DEPTH {
return Err(format!(
"query argument `{key}` nests deeper than the maximum of {}",
crate::query_string::MAX_DEPTH
));
}
if key.len() > MAX_QUERY_BYTES {
return Err(format!(
"query argument key exceeds {MAX_QUERY_BYTES} bytes"
));
}
out.visit()?;
let is_container = matches!(value, Value::Array(_) | Value::Object(_));
let before = out.pairs.len();
match value {
Value::Null => {}
Value::Array(items) if items.is_empty() => {
return Err(format!(
"query argument `{key}` is an empty array; a query string cannot express an \
empty sequence — omit the argument, or move the field to a JSON body"
));
}
Value::Array(items) => {
if let Some(position) = items.iter().position(Value::is_null) {
return Err(format!(
"query argument `{key}[{position}]` is null; a query string cannot \
express a null element — omit it, or move the field to a JSON body"
));
}
let all_scalar = items
.iter()
.all(|item| !matches!(item, Value::Array(_) | Value::Object(_)));
if all_scalar {
if let [only] = items.as_slice() {
out.push(&format!("{key}[0]"), query_scalar(only))?;
} else {
for item in items {
out.push(key, query_scalar(item))?;
}
}
} else {
for (index, item) in items.iter().enumerate() {
encode_query_arg(&format!("{key}[{index}]"), item, depth + 1, out)?;
}
}
}
Value::Object(fields) if fields.is_empty() => {
return Err(format!(
"query argument `{key}` is an empty object; a query string cannot express an \
empty object — omit the argument, or move the field to a JSON body"
));
}
Value::Object(fields)
if {
let mut emitting = fields
.iter()
.filter(|(_, value)| !value.is_null())
.peekable();
emitting.peek().is_some()
&& emitting.all(|(field, _)| {
!field.is_empty() && field.bytes().all(|b| b.is_ascii_digit())
})
} =>
{
return Err(format!(
"query argument `{key}` is an object whose field names are all numeric; a \
query string cannot tell that from a sequence — rename a field, or move the \
field to a JSON body"
));
}
Value::Object(fields) => {
for (field, nested) in fields {
if !is_expressible_field_name(field) {
return Err(inexpressible_field_name(key));
}
encode_query_arg(&format!("{key}[{field}]"), nested, depth + 1, out)?;
}
}
scalar => out.push(key, query_scalar(scalar))?,
}
if is_container && out.pairs.len() == before {
return Err(format!(
"query argument `{key}` carries no value; a query string cannot express a \
container whose every field is null — omit it, or move the field to a JSON body"
));
}
Ok(())
}
fn replace_path_param(path: &str, name: &str, value: &str) -> String {
let mut out = String::with_capacity(path.len());
let mut rest = path;
while let Some(start) = rest.find('{') {
out.push_str(&rest[..start]);
let after = &rest[start + 1..];
if let Some(end) = after.find('}') {
let inner = &after[..end];
let capture = inner.split(':').next().unwrap_or(inner).trim();
if capture == name {
out.push_str(value);
} else {
out.push('{');
out.push_str(inner);
out.push('}');
}
rest = &after[end + 1..];
} else {
out.push_str(&rest[start..]);
return out;
}
}
out.push_str(rest);
out
}
fn tool_ok(text: &str) -> Value {
json!({
"content": [ { "type": "text", "text": text } ],
"isError": false,
})
}
fn tool_error(text: &str) -> Value {
json!({
"content": [ { "type": "text", "text": text } ],
"isError": true,
})
}
fn success(id: Value, result: Value) -> Value {
let mut obj = serde_json::Map::new();
obj.insert("jsonrpc".into(), json!("2.0"));
obj.insert("id".into(), id);
obj.insert("result".into(), result);
Value::Object(obj)
}
fn error(id: Value, code: i64, message: &str) -> Value {
let mut obj = serde_json::Map::new();
obj.insert("jsonrpc".into(), json!("2.0"));
obj.insert("id".into(), id);
obj.insert("error".into(), json!({ "code": code, "message": message }));
Value::Object(obj)
}
fn error_with_data(id: Value, code: i64, message: &str, data: Value) -> Value {
let mut out = error(id, code, message);
out["error"]["data"] = data;
out
}
fn parse_error(message: &str) -> Value {
json!({
"jsonrpc": "2.0",
"id": Value::Null,
"error": { "code": -32700, "message": format!("parse error: {message}") },
})
}
fn json_response(value: &Value) -> Response {
(
[(header::CONTENT_TYPE, "application/json")],
serde_json::to_string(value).unwrap_or_else(|_| "{}".to_owned()),
)
.into_response()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openapi::{SchemaEntry, SchemaKind};
fn doc(method: &'static str, path: &'static str, op: &'static str) -> ApiDoc {
ApiDoc {
method,
path,
operation_id: op,
success_status: 200,
response: Some(SchemaEntry {
name: "Todo",
kind: SchemaKind::Ref,
identity: None,
}),
..Default::default()
}
}
#[test]
fn strip_total_metric_drops_only_total() {
assert_eq!(
strip_total_metric("total;dur=5.000, db;dur=1.250;desc=\"2 queries\""),
Some("db;dur=1.250;desc=\"2 queries\"".to_owned())
);
assert_eq!(
strip_total_metric("app;dur=1.500, total;dur=42.000"),
Some("app;dur=1.500".to_owned())
);
assert_eq!(strip_total_metric("total;dur=9.000"), None);
assert_eq!(strip_total_metric("Total"), None);
assert_eq!(strip_total_metric(" "), None);
}
#[test]
fn opt_in_required_without_hatch() {
let mut d = doc("GET", "/a", "a");
assert!(!should_expose(&d, false), "no opt-in => not exposed");
d.mcp_tool = true;
assert!(should_expose(&d, false));
}
#[test]
fn exclude_always_wins() {
let mut d = doc("GET", "/a", "a");
d.mcp_tool = true;
d.mcp_exclude = true;
assert!(!should_expose(&d, false));
assert!(!should_expose(&d, true));
}
#[test]
fn hatch_includes_reads_excludes_unopted_writes() {
let read = doc("GET", "/a", "a");
let write = doc("POST", "/a", "b");
assert!(should_expose(&read, true));
assert!(!should_expose(&write, true), "mutating verb needs opt-in");
}
#[test]
fn hatch_still_allows_opted_in_writes() {
let mut write = doc("POST", "/a", "b");
write.mcp_tool = true;
assert!(should_expose(&write, true));
}
#[test]
fn html_routes_are_ineligible() {
let mut d = doc("GET", "/page", "page");
d.response = None; d.mcp_tool = true;
assert!(!should_expose(&d, false));
}
#[test]
fn no_content_delete_with_opt_in_is_eligible() {
let mut d = doc("DELETE", "/api/widgets/{id}", "widget_api_delete");
d.response = None;
d.success_status = 204;
d.mcp_tool = true;
assert!(
should_expose(&d, false),
"opted-in 204 route is a deliberate empty success contract"
);
}
#[test]
fn reset_content_205_with_opt_in_is_eligible() {
let mut d = doc("POST", "/api/forms/clear", "clear_form");
d.response = None;
d.success_status = 205;
d.mcp_tool = true;
assert!(should_expose(&d, false));
let tools = derive_tools(&[d], false, None);
assert_eq!(tools.len(), 1, "opted-in 205 route must derive a tool");
}
const fn test_grant(
name: &'static str,
reversibility: crate::agent_authority::Reversibility,
) -> crate::agent_authority::Grant {
crate::agent_authority::Grant {
name,
writes: &[],
unbounded_writes: &[],
tenant_scope: crate::agent_authority::TenantScope::Scoped,
outbound: &[],
webhooks: &[],
jobs: &[],
rate: None,
spend: None,
reversibility,
location: "autumn/src/mcp.rs",
}
}
const fn test_authority(
action: &'static str,
grant: &'static crate::agent_authority::Grant,
) -> crate::agent_authority::AgentAuthority {
crate::agent_authority::AgentAuthority {
action,
module_path: "autumn_web::mcp::tests",
location: "autumn/src/mcp.rs",
grant,
effects: &[],
asserted_effect_free_sites: 0,
asserted_effect_free: &[],
}
}
static REVERSIBLE_GRANT: crate::agent_authority::Grant = test_grant(
"DraftOnly",
crate::agent_authority::Reversibility::Reversible,
);
static COMPENSABLE_GRANT: crate::agent_authority::Grant = test_grant(
"RefundDrafter",
crate::agent_authority::Reversibility::Compensable,
);
static IRREVERSIBLE_GRANT: crate::agent_authority::Grant = test_grant(
"PayoutSender",
crate::agent_authority::Reversibility::Irreversible,
);
static REVERSIBLE_AUTHORITY: crate::agent_authority::AgentAuthority =
test_authority("draft_note", &REVERSIBLE_GRANT);
static COMPENSABLE_AUTHORITY: crate::agent_authority::AgentAuthority =
test_authority("draft_refund", &COMPENSABLE_GRANT);
static IRREVERSIBLE_AUTHORITY: crate::agent_authority::AgentAuthority =
test_authority("send_payout", &IRREVERSIBLE_GRANT);
fn governed_doc(
method: &'static str,
op: &'static str,
authority: Option<&'static crate::agent_authority::AgentAuthority>,
) -> ApiDoc {
let mut d = doc(method, "/api/refunds", op);
d.mcp_tool = true;
d.agent_authority = authority;
d
}
#[test]
fn derive_tools_carries_the_handler_authority_onto_the_tool() {
let tools = derive_tools(
&[governed_doc(
"POST",
"draft_refund",
Some(&COMPENSABLE_AUTHORITY),
)],
false,
None,
);
assert_eq!(tools.len(), 1, "the opted-in route must derive a tool");
let authority = tools[0]
.agent_authority()
.expect("a governed handler's tool must carry its authority envelope");
assert_eq!(authority.grant.name, "RefundDrafter");
assert_eq!(
authority.grant.reversibility,
crate::agent_authority::Reversibility::Compensable
);
assert_eq!(authority.action, "draft_refund");
}
#[test]
fn derive_tools_leaves_an_ungoverned_tool_without_an_authority() {
let tools = derive_tools(&[governed_doc("POST", "draft_refund", None)], false, None);
assert_eq!(tools.len(), 1);
assert!(
tools[0].agent_authority().is_none(),
"a handler with no #[agent_operable] derives an ungoverned tool"
);
}
#[test]
fn destructive_hint_follows_declared_reversibility() {
let hint = |op, authority| {
let tools = derive_tools(&[governed_doc("POST", op, authority)], false, None);
tools[0].annotations().get("destructiveHint").cloned()
};
assert_eq!(
hint("send_payout", Some(&IRREVERSIBLE_AUTHORITY)),
Some(json!(true)),
"an irreversible action is destructive"
);
assert_eq!(
hint("draft_refund", Some(&COMPENSABLE_AUTHORITY)),
Some(json!(true)),
"a compensable action still needs a compensating step, so warn"
);
assert_eq!(
hint("draft_note", Some(&REVERSIBLE_AUTHORITY)),
Some(json!(false)),
"a POST proved to be bounded writes only is explicitly NOT destructive"
);
}
#[test]
fn destructive_hint_falls_back_to_the_verb_rule_when_ungoverned() {
let post = derive_tools(&[governed_doc("POST", "draft_refund", None)], false, None);
assert!(
post[0].annotations().get("destructiveHint").is_none(),
"an ungoverned POST keeps the verb rule's silence: {:?}",
post[0].annotations()
);
let delete = derive_tools(
&[governed_doc("DELETE", "delete_refund", None)],
false,
None,
);
assert_eq!(
delete[0].annotations().get("destructiveHint"),
Some(&json!(true)),
"DELETE stays the destructive verb when no authority is declared"
);
}
#[test]
fn the_delete_verb_is_a_floor_a_reversible_grant_cannot_clear() {
let tools = derive_tools(
&[governed_doc(
"DELETE",
"archive_note",
Some(&REVERSIBLE_AUTHORITY),
)],
false,
None,
);
assert_eq!(
tools[0].annotations().get("destructiveHint"),
Some(&json!(true)),
"one unproved adjective in a grant must not clear a DELETE's warning: {:?}",
tools[0].annotations()
);
}
fn tool_with_path_param() -> McpTool {
let mut t = tool("POST", "/api/refunds/{id}", true, false);
t.path_params = vec!["id".to_owned()];
t
}
#[test]
fn argument_names_records_the_tools_own_surface_verbatim() {
let names = argument_names(
&tool_with_path_param(),
&json!({ "body": {}, "id": "7", "query": {} }),
);
assert_eq!(names, "body,id,query", "sorted, and no unknown count");
}
#[test]
fn argument_names_counts_unrecognised_keys_instead_of_quoting_them() {
let hostile = "\n2026-09-02T00:00:00Z INFO agent tool call actor=admin";
let mut args = serde_json::Map::new();
args.insert("body".to_owned(), json!({}));
args.insert(hostile.to_owned(), json!(1));
args.insert("ssn-123-45-6789".to_owned(), json!(2));
let names = argument_names(&tool_with_path_param(), &Value::Object(args));
assert_eq!(names, "body,+2 unknown");
assert!(
!names.contains('\n') && !names.contains("ssn-"),
"no caller-chosen byte may appear: {names}"
);
}
#[test]
fn argument_names_of_an_all_unknown_call_is_only_a_count() {
let names = argument_names(
&tool_with_path_param(),
&json!({ "\u{1f648}": 1, "../../etc": 2 }),
);
assert_eq!(names, "+2 unknown");
}
#[test]
fn argument_names_addresses_a_catch_all_param_by_its_bare_name() {
let mut t = tool("GET", "/files/{*rest}", false, false);
t.path_params = vec!["*rest".to_owned()];
assert_eq!(argument_names(&t, &json!({ "rest": "a/b" })), "rest");
}
#[test]
fn argument_names_of_a_non_object_is_empty() {
assert_eq!(argument_names(&tool_with_path_param(), &json!(null)), "");
}
#[test]
fn render_effects_caps_the_list_and_marks_the_truncation() {
use crate::agent_authority::{Effect, EffectKind, EffectProvenance};
const fn effect(subject: &'static str) -> Effect {
Effect {
kind: EffectKind::Write,
subject,
location: "mcp.rs:0",
provenance: EffectProvenance::Syntactic,
}
}
const EFFECTS: [Effect; MAX_AUDITED_EFFECTS + 1] =
[effect("Refund"); MAX_AUDITED_EFFECTS + 1];
let rendered = render_effects(&EFFECTS);
assert_eq!(
rendered.matches("write:Refund").count(),
MAX_AUDITED_EFFECTS,
"an oversized effect set is capped, not spilled into every row: {rendered}"
);
assert!(
rendered.ends_with(",\u{2026}"),
"the cap must be visible in the row, not silent: {rendered}"
);
assert_eq!(render_effects(&EFFECTS[..1]), "write:Refund");
assert_eq!(render_effects(&[]), "");
}
#[test]
fn only_a_completed_stream_and_an_intact_body_count_as_success() {
assert!(Disposition::Settled.is_success());
assert!(Disposition::Stream(StreamState::Completed).is_success());
assert!(!Disposition::Stream(StreamState::Aborted).is_success());
assert!(!Disposition::Stream(StreamState::Errored).is_success());
assert!(!Disposition::Buffered(BufferedFailure::Overflow).is_success());
assert!(!Disposition::Buffered(BufferedFailure::BodyError).is_success());
assert_eq!(
Disposition::Stream(StreamState::Aborted).stream_state(),
Some("aborted")
);
assert_eq!(Disposition::Stream(StreamState::Aborted).result(), None);
assert_eq!(
Disposition::Buffered(BufferedFailure::Overflow).result(),
Some("body_overflow")
);
assert_eq!(
Disposition::Buffered(BufferedFailure::BodyError).result(),
Some("body_error")
);
assert_eq!(
Disposition::Buffered(BufferedFailure::Overflow).stream_state(),
None
);
assert_eq!(Disposition::Settled.stream_state(), None);
assert_eq!(Disposition::Settled.result(), None);
}
#[tokio::test]
async fn a_length_limit_anywhere_in_the_source_chain_is_an_overflow() {
let overflow = axum::body::to_bytes(axum::body::Body::from(vec![0_u8; 32]), 8)
.await
.expect_err("a body past the cap must not collect");
assert_eq!(
BufferedFailure::classify(&overflow),
BufferedFailure::Overflow
);
assert_eq!(BufferedFailure::Overflow.as_str(), "body_overflow");
let other = axum::Error::new(std::io::Error::other("peer went away"));
assert_eq!(
BufferedFailure::classify(&other),
BufferedFailure::BodyError
);
assert_eq!(BufferedFailure::BodyError.as_str(), "body_error");
}
#[test]
fn read_only_hint_is_unaffected_by_the_authority() {
let tools = derive_tools(
&[governed_doc(
"POST",
"draft_refund",
Some(&REVERSIBLE_AUTHORITY),
)],
false,
None,
);
assert_eq!(tools[0].annotations()["readOnlyHint"], json!(false));
}
#[test]
fn opted_in_202_is_skipped_not_exposed() {
let mut d = doc("POST", "/api/enqueue", "enqueue");
d.response = None;
d.success_status = 202;
d.mcp_tool = true;
assert!(!should_expose(&d, false));
assert!(derive_tools(&[d], false, None).is_empty());
}
#[test]
fn no_content_without_opt_in_stays_hidden() {
let mut d = doc("DELETE", "/api/widgets/{id}", "widget_api_delete");
d.response = None;
d.success_status = 204;
assert!(!should_expose(&d, false), "no opt-in => hidden");
}
#[test]
fn no_content_under_hatch_pins_behavior() {
let mut d = doc("GET", "/api/ping", "ping");
d.response = None;
d.success_status = 204;
assert!(should_expose(&d, true));
assert!(!should_expose(&d, false), "still requires hatch or opt-in");
}
#[test]
fn derive_tools_keeps_opted_in_no_content_route() {
let mut d = doc("DELETE", "/api/widgets/{id}", "widget_api_delete");
d.response = None;
d.success_status = 204;
d.mcp_tool = true;
let tools = derive_tools(&[d], false, None);
assert_eq!(tools.len(), 1, "204 opt-in must derive a tool");
assert_eq!(tools[0].name, "widget_api_delete");
assert_eq!(tools[0].annotations["destructiveHint"], true);
}
#[test]
fn streaming_tool_is_eligible_without_response_schema() {
let mut d = doc("GET", "/api/search", "search");
d.response = None;
d.mcp_stream = true;
assert!(
!should_expose(&d, false),
"stream without opt-in stays hidden"
);
d.mcp_tool = true;
assert!(
should_expose(&d, false),
"opted-in streaming tool is exposed"
);
d.mcp_exclude = true;
assert!(!should_expose(&d, false));
}
#[test]
fn streaming_get_is_included_under_the_hatch() {
let mut d = doc("GET", "/api/search", "search");
d.response = None;
d.mcp_stream = true;
assert!(should_expose(&d, true));
let mut w = doc("POST", "/api/search", "search2");
w.response = None;
w.mcp_stream = true;
assert!(!should_expose(&w, true));
}
#[test]
fn accept_header_gates_sse() {
assert!(accept_includes_event_stream(
"application/json, text/event-stream"
));
assert!(accept_includes_event_stream("text/event-stream;q=1.0"));
assert!(accept_includes_event_stream("text/*"));
assert!(!accept_includes_event_stream("application/json"));
assert!(!accept_includes_event_stream("*/*"));
}
#[test]
fn sse_parser_splits_frames_and_joins_data() {
let mut p = SseWireParser::new();
let mut events = p.push(b"data: line1\ndata: line2\n\n");
events.extend(p.push(b": keep-alive\n\nevent: result\ndata: final\n\n"));
assert_eq!(events.len(), 2);
assert_eq!(events[0].event, None);
assert_eq!(events[0].data, "line1\nline2");
assert_eq!(events[1].event.as_deref(), Some("result"));
assert_eq!(events[1].data, "final");
}
#[test]
fn sse_parser_handles_chunk_boundaries_mid_frame() {
let mut p = SseWireParser::new();
assert!(p.push(b"data: hel").is_empty());
assert!(p.push(b"lo\n").is_empty());
let events = p.push(b"\n");
assert_eq!(events.len(), 1);
assert_eq!(events[0].data, "hello");
}
#[test]
fn progress_notification_plain_and_structured() {
let token = json!("tok");
let plain = progress_notification(&token, 2.0, "working");
assert_eq!(plain["method"], "notifications/progress");
assert_eq!(plain["params"]["progressToken"], "tok");
assert_eq!(plain["params"]["progress"], 2.0);
assert_eq!(plain["params"]["message"], "working");
let structured = progress_notification(
&token,
99.0,
r#"{"progress":50,"total":100,"message":"half"}"#,
);
assert_eq!(structured["params"]["progress"], 50);
assert_eq!(structured["params"]["total"], 100);
assert_eq!(structured["params"]["message"], "half");
}
#[test]
fn collapse_sse_body_prefers_result_frames() {
let joined = collapse_sse_body(b"data: a\n\ndata: b\n\n");
assert_eq!(joined, "a\nb");
let result = collapse_sse_body(b"data: progress\n\nevent: result\ndata: done\n\n");
assert_eq!(result, "done");
}
#[test]
fn annotations_track_method() {
assert_eq!(
annotations_for("GET", "t", None)["readOnlyHint"],
json!(true)
);
assert_eq!(
annotations_for("POST", "t", None)["readOnlyHint"],
json!(false)
);
assert_eq!(
annotations_for("DELETE", "t", None)["destructiveHint"],
json!(true)
);
assert!(
annotations_for("POST", "t", None)
.get("destructiveHint")
.is_none()
);
}
#[test]
fn input_schema_includes_path_param_and_body() {
let mut d = doc("POST", "/users/{id}", "create");
d.path_params = &["id"];
d.request_body = Some(SchemaEntry {
name: "NewUser",
kind: SchemaKind::Ref,
identity: None,
});
let schema = build_input_schema(
&d,
&serde_json::Map::new(),
&crate::openapi::SchemaComponentIndex::default(),
);
assert_eq!(schema["type"], "object");
assert!(schema["properties"]["id"].is_object());
assert!(schema["properties"]["body"].is_object());
let required = schema["required"].as_array().unwrap();
assert!(required.contains(&json!("id")));
assert!(required.contains(&json!("body")));
}
#[test]
fn detect_degradation_flags_opaque_body_and_query() {
let schema = json!({
"type": "object",
"properties": {
"query": { "type": "object", "title": "Q" },
"body": { "type": "object", "title": "B" },
},
});
let degradations = detect_schema_degradations(&schema, true, true);
assert!(degradations.contains(&SchemaDegradation::OpaqueQuery));
assert!(degradations.contains(&SchemaDegradation::OpaqueBody));
}
#[test]
fn detect_degradation_ignores_field_accurate_schemas() {
let schema = json!({
"type": "object",
"properties": {
"query": {
"type": "object",
"properties": { "page": { "type": "integer" }, "tags": { "type": "array", "items": { "type": "string" } } },
},
"body": { "$ref": "#/$defs/Real" },
},
"$defs": { "Real": { "type": "object", "properties": { "a": { "type": "string" } } } },
});
assert!(detect_schema_degradations(&schema, true, true).is_empty());
}
#[test]
fn detect_degradation_accepts_structured_query_fields() {
let schema = json!({
"type": "object",
"properties": {
"query": {
"type": "object",
"properties": {
"filter": { "type": "object", "properties": { "min": { "type": "integer" } } },
"rows": { "type": "array", "items": { "type": "object", "properties": { "k": { "type": "string" } } } },
"ref_filter": { "$ref": "#/$defs/Filter" },
"maybe_filter": { "oneOf": [ { "$ref": "#/$defs/Filter" }, { "type": "null" } ] },
"maybe_rows": { "oneOf": [
{ "type": "array", "items": { "$ref": "#/$defs/Filter" } },
{ "type": "null" },
] },
},
},
},
"$defs": { "Filter": { "type": "object", "properties": { "min": { "type": "integer" } } } },
});
assert!(
detect_schema_degradations(&schema, true, false).is_empty(),
"structured query fields are honored, not degraded"
);
}
#[test]
fn an_actor_id_keeps_only_id_characters_and_is_capped() {
assert_eq!(clean_actor("user-42").as_deref(), Some("user-42"));
assert_eq!(clean_actor("a\nb\r\u{0}c").as_deref(), Some("abc"));
assert_eq!(clean_actor("\n"), None);
assert_eq!(
clean_actor("ad\u{202E}min\u{200B}").as_deref(),
Some("admin")
);
assert_eq!(clean_actor("a\u{2028}b\u{2029}c").as_deref(), Some("abc"));
assert_eq!(clean_actor("a\u{2060}b\u{00AD}c").as_deref(), Some("abc"));
assert_eq!(
clean_actor("agent:key_1.x@example.com").as_deref(),
Some("agent:key_1.x@example.com")
);
assert_eq!(clean_actor(&"x".repeat(500)).map(|a| a.len()), Some(128));
}
#[test]
fn encode_query_arg_renders_the_extractor_dialect() {
let mut pairs = QueryPairs::new();
encode_query_arg("page", &json!(2), 1, &mut pairs).expect("scalar");
encode_query_arg("tags", &json!(["a", "b"]), 1, &mut pairs).expect("scalar array");
encode_query_arg("filter", &json!({ "status": "open" }), 1, &mut pairs).expect("object");
encode_query_arg("items", &json!([{ "sku": "A" }]), 1, &mut pairs).expect("object array");
assert_eq!(
pairs.pairs,
vec![
("page".to_owned(), "2".to_owned()),
("tags".to_owned(), "a".to_owned()),
("tags".to_owned(), "b".to_owned()),
("filter[status]".to_owned(), "open".to_owned()),
("items[0][sku]".to_owned(), "A".to_owned()),
]
);
}
#[test]
fn encode_query_arg_omits_null_fields_rather_than_stringifying_them() {
let mut pairs = QueryPairs::new();
encode_query_arg("page", &json!(null), 1, &mut pairs).expect("null");
assert!(
pairs.pairs.is_empty(),
"a null argument renders no pair: {:?}",
pairs.pairs
);
let mut pairs = QueryPairs::new();
encode_query_arg(
"filter",
&json!({ "status": null, "kind": "a" }),
1,
&mut pairs,
)
.expect("null field beside a carried one");
assert_eq!(
pairs.pairs,
vec![("filter[kind]".to_owned(), "a".to_owned())]
);
}
#[test]
fn encode_query_arg_refuses_values_a_query_string_cannot_carry() {
for (label, value) in [
("empty array", json!({ "tags": [] })),
("empty object", json!({ "filter": {} })),
("empty nested array", json!({ "matrix": [[1], []] })),
("null array element", json!({ "tags": ["a", null] })),
(
"null object-array element",
json!({ "items": [null, { "sku": "A" }] }),
),
("bracket in field name", json!({ "filter": { "a[b]": 1 } })),
("empty field name", json!({ "filter": { "": 1 } })),
] {
let mut pairs = QueryPairs::new();
let Value::Object(fields) = &value else {
unreachable!()
};
let result = fields
.iter()
.try_for_each(|(key, v)| encode_query_arg(key, v, 1, &mut pairs));
assert!(
result.is_err(),
"{label} must be refused, not silently dropped"
);
}
}
#[test]
fn encode_query_arg_refuses_containers_that_collapse_to_nothing() {
let mut pairs = QueryPairs::new();
assert!(
encode_query_arg(
"items",
&json!([{ "note": null }, { "note": "x" }]),
1,
&mut pairs
)
.is_err(),
"an all-null element must not silently shorten the sequence"
);
let mut pairs = QueryPairs::new();
assert!(encode_query_arg("filter", &json!({ "note": null }), 1, &mut pairs).is_err());
let mut pairs = QueryPairs::new();
encode_query_arg("filter", &json!({ "a": 1, "b": null }), 1, &mut pairs).expect("partial");
assert_eq!(pairs.pairs, vec![("filter[a]".to_owned(), "1".to_owned())]);
}
#[test]
fn top_level_query_keys_are_validated_like_nested_field_names() {
assert!(!is_expressible_field_name("filter[x]"));
assert!(!is_expressible_field_name(""));
assert!(!is_expressible_field_name("a]b"));
assert!(is_expressible_field_name("filter"));
assert!(is_expressible_field_name("sort_by"));
}
#[test]
fn encode_query_arg_bounds_its_expansion() {
let wide: Vec<Value> = (0..(MAX_QUERY_PAIRS * 2)).map(|i| json!(i)).collect();
let mut pairs = QueryPairs::new();
assert!(encode_query_arg("tags", &Value::Array(wide), 1, &mut pairs).is_err());
let mut pairs = QueryPairs::new();
let huge = "k".repeat(MAX_QUERY_BYTES + 1);
assert!(encode_query_arg(&huge, &json!({ "a": 1 }), 1, &mut pairs).is_err());
}
#[test]
fn encode_query_arg_charges_null_leaves_for_the_traversal_they_cost() {
let nulls: serde_json::Map<String, Value> = (0..(MAX_QUERY_NODES * 2))
.map(|i| (format!("f{i}"), Value::Null))
.collect();
let mut pairs = QueryPairs::new();
let err = encode_query_arg("root", &Value::Object(nulls), 1, &mut pairs)
.expect_err("a traversal this wide must be refused");
assert!(err.contains("traverse"), "{err}");
assert!(
pairs.nodes <= MAX_QUERY_NODES + 1,
"traversal must stop at the bound, visited {}",
pairs.nodes
);
let mut pairs = QueryPairs::new();
encode_query_arg(
"filter",
&json!({ "status": "open", "tag": null }),
1,
&mut pairs,
)
.expect("a small object still encodes");
}
#[test]
fn encode_query_arg_refuses_an_all_numeric_keyed_object() {
let mut pairs = QueryPairs::new();
let err = encode_query_arg("counts", &json!({ "0": 5, "1": 6 }), 1, &mut pairs)
.expect_err("an all-numeric-keyed object is ambiguous on the wire");
assert!(err.contains("all numeric"), "{err}");
let mut pairs = QueryPairs::new();
encode_query_arg("counts", &json!({ "0": 5, "total": 6 }), 1, &mut pairs)
.expect("a mixed-key object round-trips as a map");
let mut pairs = QueryPairs::new();
assert!(encode_query_arg("filter", &json!({ "by": { "7": "x" } }), 1, &mut pairs).is_err());
let mut pairs = QueryPairs::new();
let err = encode_query_arg("counts", &json!({ "": 1 }), 1, &mut pairs)
.expect_err("an empty field name is inexpressible");
assert!(err.contains("may not be empty"), "{err}");
}
#[test]
fn a_singleton_scalar_array_keeps_its_container_kind() {
let mut pairs = QueryPairs::new();
encode_query_arg("tags", &json!(["only"]), 1, &mut pairs).expect("encodes");
assert_eq!(
pairs.pairs,
vec![("tags[0]".to_owned(), "only".to_owned())],
"a singleton must pin its kind"
);
let mut pairs = QueryPairs::new();
encode_query_arg("tags", &json!(["a", "b"]), 1, &mut pairs).expect("encodes");
assert_eq!(
pairs.pairs,
vec![
("tags".to_owned(), "a".to_owned()),
("tags".to_owned(), "b".to_owned())
]
);
}
#[test]
fn the_numeric_object_check_ignores_fields_that_emit_nothing() {
let mut pairs = QueryPairs::new();
let err = encode_query_arg("counts", &json!({ "0": 5, "kind": null }), 1, &mut pairs)
.expect_err("the null field cannot disambiguate");
assert!(err.contains("all numeric"), "{err}");
let mut pairs = QueryPairs::new();
encode_query_arg("counts", &json!({ "0": 5, "kind": "x" }), 1, &mut pairs)
.expect("an emitting named key resolves the shape");
let mut pairs = QueryPairs::new();
let err = encode_query_arg("counts", &json!({ "0": null }), 1, &mut pairs)
.expect_err("an all-null object carries no value");
assert!(err.contains("carries no value"), "{err}");
}
#[test]
fn encode_query_arg_rejects_nesting_past_the_decoder_cap() {
let mut value = json!("leaf");
for _ in 0..crate::query_string::MAX_DEPTH {
value = json!({ "deeper": value });
}
let mut pairs = QueryPairs::new();
assert!(encode_query_arg("root", &value, 1, &mut pairs).is_err());
}
#[test]
fn replace_path_param_handles_regex_captures() {
assert_eq!(replace_path_param("/u/{id}", "id", "7"), "/u/7");
assert_eq!(replace_path_param("/u/{id:[0-9]+}", "id", "7"), "/u/7");
assert_eq!(
replace_path_param("/u/{id}/p/{pid}", "pid", "9"),
"/u/{id}/p/9"
);
}
fn tool(method: &str, path: &str, has_body: bool, has_query: bool) -> McpTool {
McpTool {
name: "t".to_owned(),
description: None,
input_schema: json!({}),
annotations: json!({}),
method: method.to_owned(),
path_template: path.to_owned(),
path_params: Vec::new(),
has_body,
has_query,
streams: false,
empty_body: false,
task: false,
agent_authority: None,
core: false,
tags: Vec::new(),
}
}
#[test]
fn build_request_rejects_missing_required_body() {
let t = tool("POST", "/api/todos", true, false);
let err =
build_request(&t, &HeaderMap::new(), &json!({}), "x-csrf-token", None).unwrap_err();
assert!(err.contains("body"), "got: {err}");
}
#[test]
fn build_request_explodes_array_query_into_repeated_keys() {
let t = tool("GET", "/api/search", false, true);
let req = build_request(
&t,
&HeaderMap::new(),
&json!({ "query": { "tags": ["a", "b"], "q": "x" } }),
"x-csrf-token",
None,
)
.expect("request builds");
let query = req.uri().query().unwrap_or_default();
assert!(query.contains("tags=a"), "got: {query}");
assert!(query.contains("tags=b"), "got: {query}");
assert!(query.contains("q=x"), "got: {query}");
assert!(
!query.contains("%5B"), "array must explode, not serialize as JSON: {query}"
);
}
#[test]
fn build_request_forwards_authorization_and_cookie() {
let t = tool("GET", "/secure", false, false);
let mut headers = HeaderMap::new();
headers.insert(header::AUTHORIZATION, "Bearer tok".parse().unwrap());
headers.insert(header::COOKIE, "autumn.sid=abc".parse().unwrap());
let req =
build_request(&t, &headers, &json!({}), "x-csrf-token", None).expect("request builds");
assert_eq!(
req.headers().get(header::AUTHORIZATION).unwrap(),
"Bearer tok"
);
assert_eq!(req.headers().get(header::COOKIE).unwrap(), "autumn.sid=abc");
}
#[test]
fn build_request_forwards_csrf_token() {
let t = tool("POST", "/api/todos", true, false);
let mut headers = HeaderMap::new();
headers.insert("x-csrf-token", "csrf123".parse().unwrap());
let req = build_request(
&t,
&headers,
&json!({ "body": { "x": 1 } }),
"x-csrf-token",
None,
)
.expect("request builds");
assert_eq!(req.headers().get("x-csrf-token").unwrap(), "csrf123");
}
#[test]
fn build_request_forwards_configured_csrf_header() {
let t = tool("POST", "/api/todos", true, false);
let mut headers = HeaderMap::new();
headers.insert("x-xsrf-token", "csrf123".parse().unwrap());
let req = build_request(
&t,
&headers,
&json!({ "body": { "x": 1 } }),
"x-xsrf-token",
None,
)
.expect("request builds");
assert_eq!(req.headers().get("x-xsrf-token").unwrap(), "csrf123");
}
#[test]
fn build_request_preserves_slashes_for_catch_all_param() {
let mut t = tool("GET", "/files/{*path}", false, false);
t.path_params = vec!["*path".to_owned()];
let req = build_request(
&t,
&HeaderMap::new(),
&json!({ "path": "a/b c/d.txt" }),
"x-csrf-token",
None,
)
.expect("request builds");
assert_eq!(req.uri().path(), "/files/a/b%20c/d.txt");
}
#[test]
fn build_request_forwards_configured_tenant_header() {
let t = tool("GET", "/api/todos", false, false);
let mut headers = HeaderMap::new();
headers.insert("x-tenant-id", "acme".parse().unwrap());
let req = build_request(
&t,
&headers,
&json!({}),
"x-csrf-token",
Some("x-tenant-id"),
)
.expect("request builds");
assert_eq!(req.headers().get("x-tenant-id").unwrap(), "acme");
let req =
build_request(&t, &headers, &json!({}), "x-csrf-token", None).expect("request builds");
assert!(req.headers().get("x-tenant-id").is_none());
}
#[test]
fn build_request_rejects_non_object_query() {
let t = tool("GET", "/api/search", false, true);
for bad in [
json!({ "query": null }),
json!({ "query": "all" }),
json!({ "query": [1, 2] }),
] {
let err = build_request(&t, &HeaderMap::new(), &bad, "x-csrf-token", None).unwrap_err();
assert!(err.contains("query"), "got: {err}");
}
assert!(build_request(&t, &HeaderMap::new(), &json!({}), "x-csrf-token", None).is_ok());
}
#[test]
fn build_request_forwards_identity_and_idempotency_headers() {
let t = tool("POST", "/api/todos", true, false);
let mut headers = HeaderMap::new();
headers.insert(header::HOST, "tenant1.example.com".parse().unwrap());
headers.insert("x-forwarded-for", "203.0.113.7".parse().unwrap());
headers.insert("x-forwarded-host", "tenant1.example.com".parse().unwrap());
headers.insert("x-real-ip", "203.0.113.7".parse().unwrap());
headers.insert("idempotency-key", "abc-123".parse().unwrap());
let req = build_request(
&t,
&headers,
&json!({ "body": { "x": 1 } }),
"x-csrf-token",
None,
)
.expect("request builds");
assert_eq!(
req.headers().get(header::HOST).unwrap(),
"tenant1.example.com"
);
assert_eq!(req.headers().get("x-forwarded-for").unwrap(), "203.0.113.7");
assert_eq!(req.headers().get("x-real-ip").unwrap(), "203.0.113.7");
assert_eq!(req.headers().get("idempotency-key").unwrap(), "abc-123");
}
#[test]
fn build_request_forwards_accept_language() {
let t = tool("GET", "/api/todos", false, false);
let mut headers = HeaderMap::new();
headers.insert("accept-language", "fr-CA,fr;q=0.9".parse().unwrap());
let req =
build_request(&t, &headers, &json!({}), "x-csrf-token", None).expect("request builds");
assert_eq!(
req.headers().get("accept-language").unwrap(),
"fr-CA,fr;q=0.9"
);
}
#[test]
fn build_request_preserves_repeated_cookie_headers() {
let t = tool("POST", "/api/todos", true, false);
let mut headers = HeaderMap::new();
headers.append("cookie", "session=abc".parse().unwrap());
headers.append("cookie", "csrf=dup1".parse().unwrap());
headers.append("cookie", "csrf=dup2".parse().unwrap());
let req = build_request(
&t,
&headers,
&json!({ "body": { "x": 1 } }),
"x-csrf-token",
None,
)
.expect("request builds");
let cookies: Vec<_> = req
.headers()
.get_all("cookie")
.iter()
.map(|v| v.to_str().unwrap().to_owned())
.collect();
assert_eq!(cookies, ["session=abc", "csrf=dup1", "csrf=dup2"]);
}
#[test]
fn x402_headers_cross_the_replay_both_ways() {
let t = tool("GET", "/api/report", false, false);
let mut headers = HeaderMap::new();
headers.insert("payment-signature", "eyJ4IjoxfQ==".parse().unwrap());
let req =
build_request(&t, &headers, &json!({}), "x-csrf-token", None).expect("request builds");
assert_eq!(req.headers()["payment-signature"], "eyJ4IjoxfQ==");
let response = Response::builder()
.header(header::SET_COOKIE, "session=abc")
.header("payment-response", "receipt")
.header("payment-required", "challenge")
.header("x-other", "kept inside")
.body(Body::empty())
.unwrap();
let replayed: Vec<_> = Dispatched::inspect(&response)
.replayed
.into_iter()
.map(|(name, value)| format!("{name}: {}", value.to_str().unwrap()))
.collect();
assert_eq!(
replayed,
[
"set-cookie: session=abc",
"payment-required: challenge",
"payment-response: receipt"
]
);
}
#[tokio::test]
async fn a_failed_read_still_returns_the_payment_receipt() {
let t = tool("GET", "/api/report", false, false);
let body = Body::from_stream(futures::stream::iter([Err::<Bytes, _>(
std::io::Error::other("boom"),
)]));
let response = Response::builder()
.header("payment-response", "receipt")
.body(body)
.unwrap();
let replayed = Dispatched::inspect(&response).replayed;
let (resp, failure) = buffered_tool_response(
&t,
json!(1),
StatusCode::OK,
false,
replayed,
response,
None,
)
.await;
assert!(failure.is_some());
assert_eq!(resp.headers()["payment-response"], "receipt");
assert_eq!(resp.headers()[header::CACHE_CONTROL], "private, no-store");
assert_eq!(
resp.headers()[header::ACCESS_CONTROL_EXPOSE_HEADERS],
"PAYMENT-REQUIRED, PAYMENT-RESPONSE"
);
}
#[tokio::test]
async fn a_paid_tool_result_is_not_cached() {
let t = tool("GET", "/api/report", false, false);
let response = Response::builder()
.header("payment-response", "receipt")
.body(Body::from("{}"))
.unwrap();
let replayed = Dispatched::inspect(&response).replayed;
let (resp, _) = buffered_tool_response(
&t,
json!(1),
StatusCode::OK,
false,
replayed,
response,
None,
)
.await;
assert_eq!(resp.headers()[header::CACHE_CONTROL], "private, no-store");
let mut headers = HeaderMap::new();
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
apply_replayed(
&mut headers,
vec![(
HeaderName::from_static("payment-required"),
HeaderValue::from_static("challenge"),
)],
);
assert_eq!(
headers
.get_all(header::CACHE_CONTROL)
.iter()
.collect::<Vec<_>>(),
["private, no-store"]
);
let response = Response::builder().body(Body::from("{}")).unwrap();
let replayed = Dispatched::inspect(&response).replayed;
let (resp, _) = buffered_tool_response(
&t,
json!(1),
StatusCode::OK,
false,
replayed,
response,
None,
)
.await;
assert!(resp.headers().get(header::CACHE_CONTROL).is_none());
}
fn trusted(hosts: &[&str]) -> crate::router::TrustedHostPolicy {
let mut config = crate::config::AutumnConfig::default();
config.security.trusted_hosts.hosts = hosts.iter().map(|h| (*h).to_owned()).collect();
crate::router::TrustedHostPolicy::from_config(&config)
}
fn server(allowed_origins: Vec<String>) -> McpServer {
server_with_trusted(allowed_origins, &[])
}
fn server_with_trusted(allowed_origins: Vec<String>, hosts: &[&str]) -> McpServer {
let cors = crate::config::CorsConfig {
allowed_origins,
..crate::config::CorsConfig::default()
};
McpServer::new(
Vec::new(),
axum::Router::new(),
McpWiring {
cors,
trusted_hosts: trusted(hosts),
tenant_header: None,
csrf_header: "x-csrf-token".to_owned(),
envelope_rate_limited: false,
envelope_load_shed: false,
state: crate::AppState::for_test(),
session_cookie: "autumn.sid".to_owned(),
tasks: McpTaskConfig::default(),
instructions: None,
catalog: McpCatalog::Full,
},
)
}
#[test]
fn origin_allowlist_enforced() {
let s = server(vec!["https://ok.example".to_owned()]);
assert!(s.origin_allowed("https://ok.example", None, None));
assert!(!s.origin_allowed("https://evil.example", None, None));
assert!(!server(Vec::new()).origin_allowed("https://any.example", None, None));
assert!(server(vec!["*".to_owned()]).origin_allowed("https://any.example", None, None));
}
#[test]
fn same_origin_allowed_without_cors_allowlist() {
let s = server_with_trusted(Vec::new(), &["app.example"]);
assert!(s.origin_allowed("https://app.example", Some("app.example"), None));
assert!(s.origin_allowed("https://app.example", Some("app.example"), Some("https")));
assert!(s.origin_allowed(
"http://localhost:8080",
Some("localhost:8080"),
Some("http")
));
assert!(!s.origin_allowed("https://evil.example", Some("app.example"), None));
assert!(!s.origin_allowed("http://app.example", Some("app.example"), Some("https")));
}
#[test]
fn same_origin_normalizes_default_ports() {
let s = server_with_trusted(Vec::new(), &["app.example"]);
assert!(s.origin_allowed(
"https://app.example",
Some("app.example:443"),
Some("https")
));
assert!(s.origin_allowed(
"https://app.example:443",
Some("app.example"),
Some("https")
));
assert!(s.origin_allowed("http://app.example", Some("app.example:80"), Some("http")));
assert!(!s.origin_allowed(
"https://app.example",
Some("app.example:8443"),
Some("https")
));
assert!(!s.origin_allowed("http://app.example:443", Some("app.example"), Some("http")));
}
#[test]
fn same_origin_rejected_for_untrusted_host() {
let s = server_with_trusted(Vec::new(), &["app.example"]);
assert!(!s.origin_allowed(
"http://attacker.example",
Some("attacker.example"),
Some("http")
));
let s = server_with_trusted(vec!["http://attacker.example".to_owned()], &["app.example"]);
assert!(s.origin_allowed(
"http://attacker.example",
Some("attacker.example"),
Some("http")
));
}
#[tokio::test]
async fn options_preflight_grants_only_allowlisted_origin() {
let s = Arc::new(server_with_trusted(
vec!["https://app.example".to_owned()],
&[],
));
let mut headers = HeaderMap::new();
headers.insert(header::ORIGIN, "https://app.example".parse().unwrap());
let resp = serve_mcp_options(axum::extract::Extension(s.clone()), headers).await;
assert_eq!(resp.status(), StatusCode::NO_CONTENT);
assert_eq!(
resp.headers()
.get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
.unwrap(),
"https://app.example"
);
assert!(
resp.headers()
.get(header::ACCESS_CONTROL_ALLOW_METHODS)
.is_some()
);
let allow_headers = resp
.headers()
.get(header::ACCESS_CONTROL_ALLOW_HEADERS)
.and_then(|v| v.to_str().ok())
.unwrap()
.to_ascii_lowercase();
assert!(
allow_headers.contains("mcp-protocol-version"),
"allow-headers missing MCP-Protocol-Version: {allow_headers}"
);
assert!(
allow_headers.contains("payment-signature"),
"allow-headers missing Payment-Signature: {allow_headers}"
);
let mut headers = HeaderMap::new();
headers.insert(header::ORIGIN, "https://evil.example".parse().unwrap());
let resp = serve_mcp_options(axum::extract::Extension(s), headers).await;
assert_eq!(resp.status(), StatusCode::NO_CONTENT);
assert!(
resp.headers()
.get(header::ACCESS_CONTROL_ALLOW_ORIGIN)
.is_none()
);
}
#[test]
fn initialize_negotiates_supported_protocol_version() {
let s = server(Vec::new());
let echoed = initialize_result(&s, &json!({ "protocolVersion": "2024-11-05" }));
assert_eq!(echoed["protocolVersion"], "2024-11-05");
let fallback = initialize_result(&s, &json!({ "protocolVersion": "3999-01-01" }));
assert_eq!(fallback["protocolVersion"], DEFAULT_PROTOCOL_VERSION);
}
}