use std::collections::BTreeMap;
use std::time::{Duration, Instant};
use rmcp::model::CallToolRequestParams;
use rmcp::service::{RoleClient, RunningService};
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
use rmcp::transport::{StreamableHttpClientTransport, TokioChildProcess};
use rmcp::ServiceExt;
use serde::{Deserialize, Serialize};
use tracing::info;
use crate::error::{Error, Result};
use crate::net::{self, NetGuard};
use crate::observe::{EventKind, RunEvent};
use crate::policy::{Act, Effect, Policy};
use crate::provider::ToolSpec;
use crate::run::{refused, PendingMedia, Watch};
use crate::state::{McpEvent, PolicyEvent, Store};
pub const MCP_TOOL_PREFIX: &str = "mcp__";
const DEFAULT_TIMEOUT_SECS: u64 = 60;
fn default_timeout_secs() -> u64 {
DEFAULT_TIMEOUT_SECS
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "transport", rename_all = "snake_case")]
pub enum McpTransport {
Stdio {
command: String,
#[serde(default)]
args: Vec<String>,
#[serde(default)]
env: BTreeMap<String, String>,
},
Http {
url: String,
#[serde(default)]
headers: BTreeMap<String, String>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct McpServer {
pub id: String,
#[serde(flatten)]
pub transport: McpTransport,
#[serde(default = "default_timeout_secs")]
pub timeout_secs: u64,
}
impl McpServer {
pub fn stdio(id: impl Into<String>, command: impl Into<String>) -> Self {
Self {
id: id.into(),
transport: McpTransport::Stdio {
command: command.into(),
args: Vec::new(),
env: BTreeMap::new(),
},
timeout_secs: DEFAULT_TIMEOUT_SECS,
}
}
pub fn http(id: impl Into<String>, url: impl Into<String>) -> Self {
Self {
id: id.into(),
transport: McpTransport::Http {
url: url.into(),
headers: BTreeMap::new(),
},
timeout_secs: DEFAULT_TIMEOUT_SECS,
}
}
pub fn with_args<I, S>(mut self, args: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
if let McpTransport::Stdio { args: a, .. } = &mut self.transport {
*a = args.into_iter().map(Into::into).collect();
}
self
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout_secs = timeout.as_secs().max(1);
self
}
fn timeout(&self) -> Duration {
Duration::from_secs(self.timeout_secs.max(1))
}
}
struct Connected {
id: String,
service: RunningService<RoleClient, ()>,
timeout: Duration,
tools: Vec<ToolSpec>,
}
pub(crate) struct McpSession {
servers: Vec<Connected>,
}
impl McpSession {
pub(crate) async fn connect(
servers: &[McpServer],
policy: &Policy,
store: &Store,
run_id: i64,
watch: &Watch<'_>,
) -> Result<Self> {
let mut connected = Vec::new();
for server in servers {
let started = Instant::now();
let service = match &server.transport {
McpTransport::Stdio { command, args, env } => {
authorize_spawn(command, policy, store, run_id, watch)?;
let mut cmd = tokio::process::Command::new(command);
cmd.args(args);
for (k, v) in env {
cmd.env(k, v);
}
let transport = TokioChildProcess::new(cmd).map_err(|e| Error::Mcp {
server: server.id.clone(),
reason: format!("could not spawn {command}: {e}"),
})?;
().serve(transport).await.map_err(|e| Error::Mcp {
server: server.id.clone(),
reason: format!("handshake failed: {e}"),
})?
}
McpTransport::Http { url, headers } => {
NetGuard::new(policy)
.tracing(store, run_id, 0)
.watching(watch, 0)
.check(url)?;
let mut config = StreamableHttpClientTransportConfig::with_uri(url.clone());
let custom: std::collections::HashMap<_, _> = headers
.iter()
.filter_map(|(k, v)| {
match (
k.parse::<reqwest::header::HeaderName>(),
v.parse::<reqwest::header::HeaderValue>(),
) {
(Ok(name), Ok(value)) => Some((name, value)),
_ => None,
}
})
.collect();
if !custom.is_empty() {
config = config.custom_headers(custom);
}
let transport =
StreamableHttpClientTransport::with_client(net::http_client(), config);
().serve(transport).await.map_err(|e| Error::Mcp {
server: server.id.clone(),
reason: format!("could not connect to {url}: {e}"),
})?
}
};
let listed = tokio::time::timeout(server.timeout(), service.list_all_tools())
.await
.map_err(|_| Error::Mcp {
server: server.id.clone(),
reason: "timed out listing tools".into(),
})?
.map_err(|e| Error::Mcp {
server: server.id.clone(),
reason: format!("could not list tools: {e}"),
})?;
let tools: Vec<ToolSpec> = listed
.iter()
.map(|t| ToolSpec {
name: tool_name(&server.id, &t.name),
description: t.description.as_deref().unwrap_or_default().to_string(),
parameters: serde_json::Value::Object((*t.input_schema).clone()),
})
.collect();
let ev = McpEvent::connected(&server.id, transport_name(&server.transport))
.with_millis(started.elapsed().as_millis() as u64);
store.record_mcp(run_id, &ev)?;
announce(watch, run_id, 0, &ev);
for t in &tools {
let ev = McpEvent::discovered(&server.id, &t.name);
store.record_mcp(run_id, &ev)?;
announce(watch, run_id, 0, &ev);
}
info!(server = %server.id, tools = tools.len(), "mcp server connected");
connected.push(Connected {
id: server.id.clone(),
service,
timeout: server.timeout(),
tools,
});
}
Ok(Self { servers: connected })
}
pub(crate) fn tool_specs(&self) -> Vec<ToolSpec> {
self.servers
.iter()
.flat_map(|s| s.tools.iter().cloned())
.collect()
}
pub(crate) fn owns(&self, name: &str) -> bool {
self.servers
.iter()
.any(|s| s.tools.iter().any(|t| t.name == name))
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn call_media(
&self,
name: &str,
arguments: &serde_json::Value,
store: &Store,
run_id: i64,
step: u32,
cap: usize,
watch: &Watch<'_>,
depth: u32,
pending_media: &mut PendingMedia,
) -> Result<String> {
#[cfg(not(feature = "media"))]
let _ = pending_media;
let Some(server) = self
.servers
.iter()
.find(|s| s.tools.iter().any(|t| t.name == name))
else {
return Ok(format!("[unknown tool {name}]"));
};
let Some(bare) = bare_name(&server.id, name) else {
return Ok(format!("[unknown tool {name}]"));
};
let mut params = CallToolRequestParams::default();
params.name = bare.to_string().into();
params.arguments = arguments.as_object().cloned();
let started = Instant::now();
let outcome = tokio::time::timeout(server.timeout, server.service.call_tool(params)).await;
let millis = started.elapsed().as_millis() as u64;
let (text, ok) = match outcome {
Err(_) => (
format!("[{name} timed out after {}s]", server.timeout.as_secs()),
false,
),
Ok(Err(e)) => (format!("[{name} failed] {e}"), false),
Ok(Ok(result)) => {
let rendered = render(&result);
let failed = result.is_error.unwrap_or(false);
#[cfg(feature = "media")]
let body = {
if !failed {
pending_media.extend(rendered.images);
}
rendered.text
};
#[cfg(not(feature = "media"))]
let body = rendered;
if failed {
(format!("[{name} reported an error] {body}"), false)
} else {
(body, true)
}
}
};
let (text, truncated) = crate::tools::cap_result(text, cap);
let ev = McpEvent::called(&server.id, name, ok)
.at_step(step)
.with_millis(millis)
.with_detail(if truncated { "truncated" } else { "" });
store.record_mcp(run_id, &ev)?;
announce(watch, run_id, depth, &ev);
Ok(text)
}
pub(crate) async fn shutdown(self, store: &Store, run_id: i64, watch: &Watch<'_>) {
for s in self.servers {
let ev = McpEvent::disconnected(&s.id);
let _ = store.record_mcp(run_id, &ev);
announce(watch, run_id, 0, &ev);
let _ = s.service.cancel().await;
}
}
}
fn announce(watch: &Watch<'_>, run_id: i64, depth: u32, e: &McpEvent) {
watch.emit(RunEvent::at_depth(
run_id,
e.step,
depth,
EventKind::Mcp {
server: e.server.clone(),
tool: e.tool.clone(),
ok: e.ok,
millis: e.millis,
},
));
}
fn authorize_spawn(
command: &str,
policy: &Policy,
store: &Store,
run_id: i64,
watch: &Watch<'_>,
) -> Result<()> {
let verdict = policy.check(Act::Exec, command);
let mut ev = if verdict.effect == Effect::Allow {
PolicyEvent::decision(0, "exec", command, "allow", "policy")
} else {
PolicyEvent::refusal(0, "exec", command)
};
ev.rule = verdict.rule.clone();
ev.layer = verdict.layer.clone();
store.record_event(run_id, &ev)?;
if verdict.effect == Effect::Allow {
Ok(())
} else {
refused(watch, run_id, 0, &ev);
Err(Error::Refused {
act: "exec".into(),
target: command.to_string(),
rule: verdict.rule,
layer: verdict.layer,
})
}
}
fn tool_name(server: &str, tool: &str) -> String {
format!("{MCP_TOOL_PREFIX}{server}__{tool}")
}
fn bare_name<'a>(server: &str, namespaced: &'a str) -> Option<&'a str> {
namespaced.strip_prefix(&format!("{MCP_TOOL_PREFIX}{server}__"))
}
fn transport_name(t: &McpTransport) -> &'static str {
match t {
McpTransport::Stdio { .. } => "stdio",
McpTransport::Http { .. } => "http",
}
}
#[cfg(feature = "media")]
pub(crate) struct Rendered {
pub text: String,
pub images: Vec<crate::provider::Media>,
}
#[cfg(feature = "media")]
pub(crate) fn render(result: &rmcp::model::CallToolResult) -> Rendered {
use rmcp::model::ContentBlock;
let mut parts: Vec<String> = Vec::new();
let mut images = Vec::new();
let mut saw_text = false;
for c in &result.content {
match c {
ContentBlock::Text(t) => {
saw_text = true;
parts.push(t.text.clone());
}
ContentBlock::Image(i) => parts.push(take_image(i, &mut images)),
ContentBlock::Audio(a) => parts.push(format!(
"[audio: {}, not attached — only images are passed to the model]",
a.mime_type
)),
ContentBlock::Resource(_) => parts.push("[embedded resource, not attached]".into()),
ContentBlock::ResourceLink(_) => parts.push("[resource link, not attached]".into()),
_ => parts.push("[non-text content, not attached]".into()),
}
}
if !saw_text {
if let Some(structured) = &result.structured_content {
parts.push(structured.to_string());
}
}
let text = if parts.is_empty() {
"(no text content)".to_string()
} else {
parts.join("\n")
};
Rendered { text, images }
}
#[cfg(feature = "media")]
fn take_image(img: &rmcp::model::ImageContent, images: &mut Vec<crate::provider::Media>) -> String {
use crate::provider::{Media, IMAGE_MEDIA_TYPES, MAX_IMAGE_BYTES};
if !IMAGE_MEDIA_TYPES.contains(&img.mime_type.as_str()) {
return format!(
"[image not attached: unsupported media type {:?}; expected one of {}]",
img.mime_type,
IMAGE_MEDIA_TYPES.join(", ")
);
}
if img.data.len() < 4 || !img.data.len().is_multiple_of(4) {
return "[image not attached: malformed base64 payload]".to_string();
}
let media = Media {
media_type: img.mime_type.clone(),
base64: img.data.clone(),
};
let bytes = media.byte_len();
if bytes > MAX_IMAGE_BYTES {
return format!(
"[image not attached: {bytes} bytes, over the {MAX_IMAGE_BYTES}-byte per-image bound]"
);
}
let line = format!("[image: {}, {bytes} bytes]", media.media_type);
images.push(media);
line
}
#[cfg(not(feature = "media"))]
fn render(result: &rmcp::model::CallToolResult) -> String {
let mut parts: Vec<String> = result
.content
.iter()
.filter_map(|c| c.as_text().map(|t| t.text.clone()))
.collect();
if parts.is_empty() {
if let Some(structured) = &result.structured_content {
parts.push(structured.to_string());
}
}
if parts.is_empty() {
return "(no text content)".to_string();
}
parts.join("\n")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tool_names_are_namespaced_and_reversible() {
let n = tool_name("files", "write_file");
assert_eq!(n, "mcp__files__write_file");
assert!(n.starts_with(MCP_TOOL_PREFIX));
assert_ne!(n, "write_file", "a server must not shadow a built-in");
assert_eq!(bare_name("files", &n), Some("write_file"));
assert_eq!(bare_name("other", &n), None);
}
#[test]
fn a_server_config_round_trips_through_serde() {
for server in [
McpServer::stdio("files", "mcp-files").with_args(["--root", "/tmp"]),
McpServer::http("remote", "https://mcp.example.com/mcp")
.with_timeout(Duration::from_secs(5)),
] {
let json = serde_json::to_string(&server).unwrap();
let back: McpServer = serde_json::from_str(&json).unwrap();
assert_eq!(server, back, "{json}");
}
}
#[test]
fn a_stdio_config_omitting_optional_fields_still_parses() {
let s: McpServer =
serde_json::from_str(r#"{"id":"files","transport":"stdio","command":"mcp-files"}"#)
.unwrap();
assert_eq!(s.timeout_secs, DEFAULT_TIMEOUT_SECS);
assert!(matches!(s.transport, McpTransport::Stdio { .. }));
}
#[test]
fn an_oversized_result_is_cut_on_a_char_boundary_and_says_so() {
let cap_chars =
crate::context::entry_cap_chars(crate::context::ContextBudget::default().max_tokens);
let (short, cut) = crate::tools::cap_result("hello".into(), cap_chars);
assert_eq!((short.as_str(), cut), ("hello", false));
let (long, cut) = crate::tools::cap_result("é".repeat(cap_chars), cap_chars);
assert!(cut);
assert!(long.contains("[truncated at"));
assert!(long.len() < 2 * cap_chars);
}
#[test]
fn a_text_only_result_renders_exactly_its_text_and_attaches_nothing() {
use rmcp::model::{CallToolResult, ContentBlock};
let r = CallToolResult::success(vec![
ContentBlock::text("first line"),
ContentBlock::text("second line"),
]);
let out = render(&r);
#[cfg(feature = "media")]
{
assert_eq!(out.text, "first line\nsecond line");
assert!(out.images.is_empty(), "text carries no images");
}
#[cfg(not(feature = "media"))]
assert_eq!(out, "first line\nsecond line");
}
#[cfg(feature = "media")]
mod media {
use super::super::{render, take_image};
use crate::provider::MAX_IMAGE_BYTES;
use rmcp::model::{CallToolResult, ContentBlock};
const PIXEL: &str = "aGVsbG8h";
fn b64(decoded_bytes: usize) -> String {
"A".repeat(decoded_bytes.div_ceil(3) * 4)
}
#[test]
fn an_image_result_attaches_the_image_and_names_it_in_the_text() {
let r = CallToolResult::success(vec![ContentBlock::image(PIXEL, "image/png")]);
let out = render(&r);
assert_eq!(out.images.len(), 1, "the image is passed through");
assert_eq!(out.images[0].media_type, "image/png");
assert_eq!(
out.images[0].base64, PIXEL,
"base64 is moved, not re-encoded"
);
assert_eq!(out.text, "[image: image/png, 6 bytes]");
assert_ne!(
out.text, "(no text content)",
"the 0.8 bug, not reintroduced"
);
}
#[test]
fn an_unsupported_media_type_is_a_note_not_an_attachment() {
let bad = CallToolResult::success(vec![ContentBlock::image(PIXEL, "image/tiff")]);
let out = render(&bad);
assert!(out.images.is_empty(), "no vendor accepts image/tiff");
assert!(
out.text.contains("unsupported media type") && out.text.contains("image/tiff"),
"{}",
out.text
);
let good = CallToolResult::success(vec![ContentBlock::image(PIXEL, "image/webp")]);
let out = render(&good);
assert_eq!(out.images.len(), 1);
assert_eq!(out.images[0].media_type, "image/webp");
}
#[test]
fn an_oversized_image_is_a_note_not_an_attachment() {
let over = CallToolResult::success(vec![ContentBlock::image(
b64(MAX_IMAGE_BYTES + 3),
"image/jpeg",
)]);
let out = render(&over);
assert!(out.images.is_empty(), "over the per-image bound");
assert!(
out.text.contains("over the") && out.text.contains("per-image bound"),
"{}",
out.text
);
let under = CallToolResult::success(vec![ContentBlock::image(
b64(MAX_IMAGE_BYTES - 3),
"image/jpeg",
)]);
let out = render(&under);
assert_eq!(out.images.len(), 1);
assert!(out.images[0].byte_len() <= MAX_IMAGE_BYTES);
}
#[test]
fn text_and_an_image_together_yield_both() {
let r = CallToolResult::success(vec![
ContentBlock::text("here is the chart"),
ContentBlock::image(PIXEL, "image/gif"),
]);
let out = render(&r);
assert_eq!(out.images.len(), 1);
assert_eq!(out.text, "here is the chart\n[image: image/gif, 6 bytes]");
}
#[test]
fn a_malformed_base64_payload_is_a_note_rather_than_a_panic() {
let mut images = Vec::new();
let note = take_image(
&rmcp::model::ImageContent::new("=", "image/png"),
&mut images,
);
assert!(images.is_empty());
assert!(note.contains("malformed base64"), "{note}");
let ok = take_image(
&rmcp::model::ImageContent::new(PIXEL, "image/png"),
&mut images,
);
assert_eq!(images.len(), 1);
assert!(ok.starts_with("[image: image/png"), "{ok}");
}
#[test]
fn non_image_content_is_described_rather_than_dropped() {
let r = CallToolResult::success(vec![ContentBlock::audio("AAAA", "audio/wav")]);
let out = render(&r);
assert!(out.images.is_empty(), "audio is not sent to any provider");
assert!(out.text.contains("audio/wav"), "{}", out.text);
assert_ne!(out.text, "(no text content)");
}
}
}