diff --git a/.aether/settings.json b/.aether/settings.json index b7de30cd8..e65f839ad 100644 --- a/.aether/settings.json +++ b/.aether/settings.json @@ -31,8 +31,7 @@ "${HOME}/.agents/rules", "--rules-dir", "${WORKSPACE}/.agents/rules" - ], - "input": null + ] } } } diff --git a/.aether/skills/rust/error-handling.md b/.aether/skills/rust/error-handling.md index db962f989..255bbea5b 100644 --- a/.aether/skills/rust/error-handling.md +++ b/.aether/skills/rust/error-handling.md @@ -19,8 +19,6 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum AgentError { - #[error("MCP error: {0}")] - McpError(#[from] mcp_utils::client::McpError), #[error("LLM error: {0}")] LlmError(#[from] llm::LlmError), #[error("IO error: {0}")] diff --git a/Cargo.lock b/Cargo.lock index bcbb36f86..6615855dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -89,7 +89,6 @@ dependencies = [ "futures", "libc", "ratatui", - "rmcp", "rustls", "schemars 1.2.2", "serde", @@ -123,6 +122,7 @@ dependencies = [ "clap", "futures", "glob", + "http", "noyalib", "reqwest 0.13.5", "rmcp", @@ -134,6 +134,7 @@ dependencies = [ "thiserror 2.0.21", "tokio", "tokio-stream", + "tokio-util", "tracing", "tracing-subscriber", "uuid", @@ -332,24 +333,19 @@ name = "aether-mcp-utils" version = "0.7.13" dependencies = [ "aether-auth", - "aether-doctest", - "aether-llm", - "aether-utils", "async-stream", "axum", "chrono", "futures", "reqwest 0.13.5", "rmcp", - "schemars 1.2.2", "serde", "serde_json", - "tempfile", "thiserror 2.0.21", "tokio", + "tokio-stream", "tokio-util", "tracing", - "tracing-subscriber", "uuid", ] @@ -360,7 +356,6 @@ dependencies = [ "aether-agent-core", "aether-doctest", "aether-llm", - "aether-mcp-utils", "aether-project", "aether-utils", "globset", @@ -6618,11 +6613,14 @@ dependencies = [ "chrono", "futures", "http", + "http-body", + "http-body-util", "indexmap 2.14.2", "oauth2", "pastey", "pin-project-lite", "process-wrap", + "rand 0.10.3", "reqwest 0.13.5", "rmcp-macros", "schemars 1.2.2", @@ -6633,6 +6631,7 @@ dependencies = [ "tokio", "tokio-stream", "tokio-util", + "tower-service", "tracing", "url", "uuid", @@ -7928,6 +7927,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] diff --git a/README.md b/README.md index f6bb9b652..05b3c9fc3 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,7 @@ Use `aether-agent-core` as a Rust library to build your own agent in ~25 lines. mcp::mcp, }; use llm::providers::anthropic::AnthropicProvider; + use aether_core::mcp::McpConfig; use std::io::{self, Write}; #[tokio::main] @@ -143,14 +144,15 @@ Use `aether-agent-core` as a Rust library to build your own agent in ~25 lines. // 1. Create a provider (reads ANTHROPIC_API_KEY from env) let llm = AnthropicProvider::new(None); - // 2. Spawn MCP tool servers from one or more mcp.json files - let mut mcp_runtime = mcp(".").from_json_files(&["mcp.json"])?.spawn().await?; - let snapshot = mcp_runtime.block_until_ready().await.ok_or("MCP bootstrap aborted")?; + // 2. Spawn MCP tool servers from an mcp.json file + let config = McpConfig::from_json_file("mcp.json")?; + let mcp = mcp(".").with_config(config)?.spawn()?; + mcp.ready().await; // 3. Build and spawn the agent let (tx, mut rx, _handle) = agent(llm) .system_prompt(Prompt::file("AGENTS.md", ".")) - .tools(mcp_runtime.handle().clone(), snapshot.tool_definitions()) + .mcp(mcp) .spawn() .await?; diff --git a/crates/aether-cli/Cargo.toml b/crates/aether-cli/Cargo.toml index 785201fb3..7df0ccb4d 100644 --- a/crates/aether-cli/Cargo.toml +++ b/crates/aether-cli/Cargo.toml @@ -29,7 +29,7 @@ name = "aether" path = "src/main.rs" [features] -testing = ["acp_utils/testing", "mcp_utils/testing"] +testing = ["acp_utils/testing", "mcp_utils/testing", "aether-core/testing"] [dependencies] aether-auth = { path = "../aether-auth", version = "0.1.17", features = ["keyring"] } @@ -40,7 +40,7 @@ aether-project = { path = "../aether-project", version = "0.5.58" } aether-telemetry = { path = "../aether-telemetry", version = "0.3.13" } llm = { package = "aether-llm", path = "../llm", features = ["bedrock", "codex"], version = "0.9.13" } mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", version = "0.7.13" } -mcp-servers = { package = "aether-mcp-servers", path = "../mcp-servers", features = ["all"], version = "0.8.4" } +mcp-servers = { package = "aether-mcp-servers", path = "../mcp-servers", features = ["subagents"], version = "0.8.4" } wisp = { package = "aether-wisp", path = "../wisp", version = "0.7.14" } ratatui = { workspace = true } acp_utils = { package = "aether-acp-utils", path = "../acp-utils", features = ["websocket"], version = "0.5.13" } @@ -49,7 +49,6 @@ clankerdiff-core = { workspace = true } clankerdiff-protocol = { workspace = true } clankerdiff-server = { workspace = true } aether-lspd = { path = "../aether-lspd", version = "0.3.3" } -rmcp = { workspace = true, features = ["client", "elicitation", "server", "transport-streamable-http-client-reqwest"] } thiserror = { workspace = true } tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros", "sync", "io-util", "net", "fs", "time", "signal", "process"] } tokio-tungstenite = { workspace = true } diff --git a/crates/aether-cli/README.md b/crates/aether-cli/README.md index 878a11047..a0e5270df 100644 --- a/crates/aether-cli/README.md +++ b/crates/aether-cli/README.md @@ -153,7 +153,7 @@ The `mcp.json` file configures MCP tool servers: - **coding** — Filesystem tools (read, write, bash, etc.) plus optional auto-read rules from configured `--rules-dir` paths - **skills** — Slash commands and reusable prompts loaded from the configured `--dir` paths -Set `"deferTools": true` on a server to keep its schemas out of the model-visible tool definitions. Use `aether mcp --help` to discover deferred servers, then `aether mcp --help` and `aether mcp --help` to inspect tools and schemas. Invoke a tool with `--json '{...}'` or a JSON object on stdin. To partition tools, use `"deferTools": { "include": [...], "exclude": [...] }`; patterns are exact MCP-local names or trailing-`*` prefixes, and `exclude` wins. +Set `"deferTools": true` on a server to keep its schemas out of the model-visible tool definitions. Use `aether mcp --help` to discover deferred servers, then `aether mcp --help` and `aether mcp --help` to inspect tools and schemas. Invoke a tool with `--json '{...}'` or a JSON object on stdin. To partition tools, use `"deferTools": { "include": [...], "exclude": [...] }`; entries are exact MCP-local names, trailing-`*` prefixes, or annotation matchers such as `{ "readOnly": true }`, and `exclude` wins. ## Slash Commands diff --git a/crates/aether-cli/src/acp/fake_prompt_mcp.rs b/crates/aether-cli/src/acp/fake_prompt_mcp.rs deleted file mode 100644 index 2e2a379fb..000000000 --- a/crates/aether-cli/src/acp/fake_prompt_mcp.rs +++ /dev/null @@ -1,83 +0,0 @@ -use rmcp::model::{ - GetPromptRequestParams, GetPromptResponse, Implementation, ListPromptsResult, PaginatedRequestParams, - Prompt as McpPrompt, PromptMessage, Role, ServerCapabilities, ServerConfig, -}; -use rmcp::service::{DynService, RequestContext}; -use rmcp::{ErrorData as McpError, RoleServer, ServerHandler}; -use std::sync::Arc; -use tokio::sync::{Notify, mpsc}; - -#[derive(Clone)] -pub(crate) struct FakePromptMcp { - prompt_name: String, - gate: Option<(Arc, Arc)>, - elicitation_results: Option>, -} - -impl FakePromptMcp { - pub(crate) fn new(prompt_name: &str) -> Self { - Self { prompt_name: prompt_name.to_string(), gate: None, elicitation_results: None } - } - - pub(crate) fn with_gate(mut self, gate: Option<(Arc, Arc)>) -> Self { - self.gate = gate; - self - } - - pub(crate) fn with_elicitation( - mut self, - results: Option>, - ) -> Self { - self.elicitation_results = results; - self - } - - pub(crate) fn into_dyn(self) -> Box> { - Box::new(self) - } -} - -impl ServerHandler for FakePromptMcp { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_prompts().build()) - .with_server_info(Implementation::new("fake-prompt-mcp", "0.1.0")) - .with_instructions("Fake MCP server exposing a single prompt for tests") - } - - fn list_prompts( - &self, - _request: Option, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - let prompt = McpPrompt::new(&self.prompt_name, Some(format!("{} command", self.prompt_name)), None); - std::future::ready(Ok(ListPromptsResult::with_all_items(vec![prompt]))) - } - - fn get_prompt( - &self, - request: GetPromptRequestParams, - context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - let result = if request.name.as_str() == self.prompt_name { - let messages = vec![PromptMessage::new_text(Role::User, format!("expanded {}", self.prompt_name))]; - Ok(rmcp::model::GetPromptResult::new(messages).into()) - } else { - Err(McpError::invalid_params(format!("Prompt '{}' not found", request.name), None)) - }; - async move { - if let Some((started, release)) = &self.gate { - started.notify_one(); - release.notified().await; - } - if let Some(results) = &self.elicitation_results { - let params = serde_json::from_value(serde_json::json!({ - "mode": "form", "message": "Continue?", "requestedSchema": {"type": "object", "properties": {}} - })) - .expect("valid elicitation request"); - let response = context.peer.create_elicitation(params).await.expect("elicitation response"); - let _ = results.send(response); - } - result - } - } -} diff --git a/crates/aether-cli/src/acp/mod.rs b/crates/aether-cli/src/acp/mod.rs index 04e93f12d..bea5b6a76 100644 --- a/crates/aether-cli/src/acp/mod.rs +++ b/crates/aether-cli/src/acp/mod.rs @@ -1,6 +1,4 @@ pub(crate) mod agent; -#[cfg(any(test, feature = "testing"))] -pub(crate) mod fake_prompt_mcp; pub(crate) mod protocol; pub mod server; pub(crate) mod session; @@ -8,7 +6,6 @@ pub(crate) mod state; #[cfg(any(test, feature = "testing"))] pub mod testing; -pub use protocol::map_mcp_prompt_to_available_command; pub use session::hooks::{IdleHook, SessionHooks}; use crate::acp::state::{AcpState, AcpStateConfig}; diff --git a/crates/aether-cli/src/acp/protocol/commands.rs b/crates/aether-cli/src/acp/protocol/commands.rs index 7b85043f7..e2cef8b6d 100644 --- a/crates/aether-cli/src/acp/protocol/commands.rs +++ b/crates/aether-cli/src/acp/protocol/commands.rs @@ -1,22 +1,17 @@ use agent_client_protocol::schema::v2::{self as acp, AvailableCommandInput, TextCommandInput}; -use rmcp::model::Prompt as McpPrompt; +use mcp_utils::gateway::ServerPrompt; -/// Converts an MCP Prompt to an ACP `AvailableCommand` -/// -/// Strips the MCP namespace from the prompt name (e.g., "`coding__web`" -> "web") -/// and creates a slash command that clients can invoke. -pub fn map_mcp_prompt_to_available_command(prompt: &McpPrompt) -> acp::AvailableCommand { - let command_name = prompt.name.split("__").last().unwrap_or(prompt.name.as_ref()).to_string(); +pub(crate) fn available_commands(prompts: &[ServerPrompt]) -> Vec { + prompts.iter().map(|listed| available_command(&listed.prompt)).collect() +} +fn available_command(prompt: &mcp_utils::model::Prompt) -> acp::AvailableCommand { let hint = prompt .arguments .as_ref() .and_then(|args| args.iter().find(|a| a.name.as_str() == "ARGUMENTS").and_then(|a| a.description.as_deref())) .unwrap_or("optional arguments"); - - let input = Some(AvailableCommandInput::Text(TextCommandInput::new(hint))); - let description = prompt.description.clone().unwrap_or_else(|| "No description available".to_string()); - - acp::AvailableCommand::new(command_name, description).input(input) + acp::AvailableCommand::new(&prompt.name, description) + .input(Some(AvailableCommandInput::Text(TextCommandInput::new(hint)))) } diff --git a/crates/aether-cli/src/acp/protocol/mcp.rs b/crates/aether-cli/src/acp/protocol/mcp.rs index 2cdc6f877..ca3e0f4f7 100644 --- a/crates/aether-cli/src/acp/protocol/mcp.rs +++ b/crates/aether-cli/src/acp/protocol/mcp.rs @@ -1,13 +1,14 @@ +use aether_core::mcp::http_transport; use agent_client_protocol::schema::v2::{HttpHeader, McpServer}; -use mcp_utils::client::{McpServer as RuntimeMcpServer, McpTransport, ToolExposure}; -use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use mcp_utils::client::{HeaderMap, Transport}; +use mcp_utils::gateway::ServerSpec; +use std::path::Path; -/// Maps ACP MCP server definitions to internal MCP servers, skipping unsupported transports or invalid headers. -pub fn map_acp_mcp_servers(servers: Vec) -> Vec { +pub fn map_acp_mcp_servers(servers: Vec, cwd: &Path) -> Vec { servers .into_iter() .filter_map(|s| { - try_map_mcp_server(s).or_else(|| { + try_map_mcp_server(s, cwd).or_else(|| { tracing::warn!("Unsupported ACP MCP transport or invalid HTTP headers, skipping server"); None }) @@ -15,36 +16,26 @@ pub fn map_acp_mcp_servers(servers: Vec) -> Vec { .collect() } -fn try_map_mcp_server(server: McpServer) -> Option { +fn try_map_mcp_server(server: McpServer, cwd: &Path) -> Option { use McpServer::{Http, Stdio}; - match server { - Stdio(stdio) => Some(RuntimeMcpServer::new( + let (name, transport) = match server { + Stdio(stdio) => ( stdio.name, - McpTransport::Stdio { + Transport::Stdio { command: stdio.command.0.to_string_lossy().into_owned(), args: stdio.args, env: stdio.env.into_iter().map(|e| (e.name, e.value)).collect(), + cwd: Some(cwd.to_path_buf()), }, - ToolExposure::ModelVisible, - )), - - Http(http) => Some(RuntimeMcpServer::new( - http.name, - McpTransport::Http(http_config(http.url, &http.headers)?.into()), - ToolExposure::ModelVisible, - )), - - _ => None, - } + ), + Http(http) => (http.name, http_transport(http.url, header_map(&http.headers)?, None)), + _ => return None, + }; + Some(ServerSpec { name, transport }) } -fn http_config(url: String, headers: &[HttpHeader]) -> Option { - let mut config = StreamableHttpClientTransportConfig::with_uri(url); - for header in headers { - // ACP supplies complete header values; rmcp's auth_header would prepend `Bearer `. - config.custom_headers.insert(header.name.parse().ok()?, header.value.parse().ok()?); - } - Some(config) +fn header_map(headers: &[HttpHeader]) -> Option { + headers.iter().map(|header| Some((header.name.parse().ok()?, header.value.parse().ok()?))).collect() } #[cfg(test)] @@ -52,6 +43,10 @@ mod tests { use super::*; use agent_client_protocol::schema::v2 as acp; + fn map(servers: Vec) -> Vec { + map_acp_mcp_servers(servers, Path::new("/workspace")) + } + #[test] fn test_map_acp_stdio_server() { let server = acp::McpServer::Stdio( @@ -60,15 +55,16 @@ mod tests { .env(vec![acp::EnvVariable::new("FOO", "bar")]), ); - let configs = map_acp_mcp_servers(vec![server]); + let configs = map(vec![server]); assert_eq!(configs.len(), 1); match &configs[0].transport { - McpTransport::Stdio { command, args, env } => { + Transport::Stdio { command, args, env, cwd } => { assert_eq!(configs[0].name, "my-server"); assert_eq!(command, "/usr/bin/server"); assert_eq!(args, &["--port", "8080"]); assert_eq!(env.get("FOO").unwrap(), "bar"); + assert_eq!(cwd.as_deref(), Some(Path::new("/workspace"))); } other => panic!("Expected Stdio, got {other:?}"), } @@ -81,26 +77,15 @@ mod tests { .headers(vec![acp::HttpHeader::new("Authorization", "Bearer token123")]), ); - let configs = map_acp_mcp_servers(vec![server]); + let configs = map(vec![server]); assert_eq!(configs.len(), 1); match &configs[0].transport { - McpTransport::Http(config) => { + Transport::Http { url, headers, oauth } => { assert_eq!(configs[0].name, "http-server"); - assert_eq!(config.transport.uri.as_ref(), "https://example.com/mcp"); - assert!(config.transport.auth_header.is_none()); - assert_eq!( - config - .transport - .custom_headers - .iter() - .find(|(name, _)| name.as_str() == "authorization") - .unwrap() - .1 - .to_str() - .unwrap(), - "Bearer token123" - ); + assert_eq!(url, "https://example.com/mcp"); + assert_eq!(headers["authorization"], "Bearer token123"); + assert!(oauth.is_none(), "an Authorization header replaces OAuth"); } other => panic!("Expected Http, got {other:?}"), } @@ -114,16 +99,9 @@ mod tests { acp::HttpHeader::new("Authorization", input), acp::HttpHeader::new("X-API-Key", "secret"), ])); - let configs = map_acp_mcp_servers(vec![server]); + let configs = map(vec![server]); match &configs[0].transport { - McpTransport::Http(config) => { - assert!(config.transport.auth_header.is_none()); - let headers: std::collections::BTreeMap<_, _> = config - .transport - .custom_headers - .iter() - .map(|(name, value)| (name.as_str(), value.to_str().unwrap())) - .collect(); + Transport::Http { headers, .. } => { assert_eq!(headers["authorization"], input); assert_eq!(headers["x-api-key"], "secret"); } @@ -138,7 +116,7 @@ mod tests { acp::McpServerHttp::new("invalid", "https://example.com/mcp") .headers(vec![acp::HttpHeader::new("Authorization", "invalid\nvalue")]), ); - assert!(map_acp_mcp_servers(vec![server]).is_empty()); + assert!(map(vec![server]).is_empty()); } #[test] @@ -148,11 +126,12 @@ mod tests { {"type": "http", "name": "remote", "url": "https://example.com/mcp"} ])) .unwrap(); - let configs = map_acp_mcp_servers(servers); - let McpTransport::Stdio { args, env, .. } = &configs[0].transport else { panic!("expected stdio") }; + let configs = map(servers); + let Transport::Stdio { args, env, .. } = &configs[0].transport else { panic!("expected stdio") }; assert!(args.is_empty() && env.is_empty()); - let McpTransport::Http(config) = &configs[1].transport else { panic!("expected http") }; - assert!(config.transport.auth_header.is_none() && config.transport.custom_headers.is_empty()); + let Transport::Http { headers, oauth, .. } = &configs[1].transport else { panic!("expected http") }; + assert!(headers.is_empty()); + assert!(oauth.is_some(), "servers without credentials can sign in with OAuth"); } #[test] @@ -161,6 +140,6 @@ mod tests { "type": "sse", "name": "unsupported", "url": "https://example.com/sse" })) .unwrap(); - assert!(map_acp_mcp_servers(vec![server]).is_empty()); + assert!(map(vec![server]).is_empty()); } } diff --git a/crates/aether-cli/src/acp/protocol/mod.rs b/crates/aether-cli/src/acp/protocol/mod.rs index 29a77f7fb..164c4810a 100644 --- a/crates/aether-cli/src/acp/protocol/mod.rs +++ b/crates/aether-cli/src/acp/protocol/mod.rs @@ -10,8 +10,6 @@ pub(crate) mod events; pub(crate) mod mcp; pub(crate) mod replay; -pub use commands::map_mcp_prompt_to_available_command; - pub(crate) fn notify( connection: &agent_client_protocol::ConnectionTo, notification: impl agent_client_protocol::JsonRpcNotification, diff --git a/crates/aether-cli/src/acp/session/actor.rs b/crates/aether-cli/src/acp/session/actor.rs index 61ffda899..12e3a1579 100644 --- a/crates/aether-cli/src/acp/session/actor.rs +++ b/crates/aether-cli/src/acp/session/actor.rs @@ -4,7 +4,7 @@ use acp_utils::elicitation; use acp_utils::notifications::{GitDiffEventPayload, McpNotification}; use aether_auth::OAuthCredentialStorage; use aether_core::events::{AgentCommand, AgentEvent, Command, TurnEvent, TurnOutcome}; -use aether_core::mcp::McpHandle; +use aether_core::mcp::McpCatalog; use aether_sessions::model::{SessionControlEvent, SessionEvent, last_session_usage}; use aether_sessions::transcript::conversation_messages_from_events; use agent_client_protocol::schema::v2::{ @@ -14,11 +14,11 @@ use agent_client_protocol::schema::v2::{ use agent_client_protocol::{Client, ConnectionTo, Error, Responder}; use clankerdiff_protocol::client::ClientCommand; use clankerdiff_protocol::shared::{DocumentUpdate, Event}; -use futures::{FutureExt, StreamExt, future::BoxFuture, stream::FuturesUnordered}; +use futures::{FutureExt, StreamExt, future::BoxFuture}; use llm::catalog::LlmModel; use llm::parser::ModelProviderParser; use llm::{ChatMessage, ContentBlock, MessageId, ProviderConnectionOverrides}; -use mcp_utils::client::{ElicitationRequest, McpClientEvent, cancel_result}; +use mcp_utils::client::{Elicitation, ElicitationRequest}; use std::collections::HashMap; use std::path::PathBuf; use std::process::Stdio; @@ -29,7 +29,7 @@ use tokio::task::{JoinHandle, JoinSet}; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tracing::{error, info, warn}; -use utils::mcp_status::McpServerStatusEntry; +use utils::mcp_status::{McpServerStatus, McpServerStatusEntry}; use super::agent_key::AgentKey; use super::agents::SessionAgents; @@ -40,17 +40,17 @@ use super::hooks::SessionHooks; use super::model::Modes; use super::prompts::{PromptQueue, SubmittedPrompt}; use super::runtime::{AgentRuntime, RuntimeFactory}; -use super::slash_commands::{expand_slash_command_in_content, send_available_commands}; -use crate::acp::protocol::commands::map_mcp_prompt_to_available_command; +use crate::acp::protocol::commands::available_commands; use crate::acp::protocol::events::{NotificationMode, project_agent_event}; use crate::acp::protocol::replay::replay_to_client; use crate::acp::state::validate_prompt_support; use crate::output::print_message; -use crate::slash_commands::dedupe_commands_by_name; +use crate::slash_commands::{expand_slash_command_in_content, prompt_commands}; use aether_sessions::SessionStore; /// Capacity of the per-session command channel feeding the actor loop. const SESSION_COMMAND_CHANNEL_CAPACITY: usize = 50; +const ELICITATION_CHANNEL_CAPACITY: usize = 32; /// A command routed to a single session's actor. The actor is the only consumer /// of these, so per-session state never needs an additional lock. @@ -200,30 +200,25 @@ pub(crate) struct SessionActor { specs: SessionAgents, runtimes: HashMap, runtime_factory: Arc, + elicitation_tx: mpsc::Sender, + elicitation_rx: mpsc::Receiver, transcript: Vec, config: SessionConfigState, modes: Modes, prompts: PromptQueue, - command_refresh: Option, SessionError>>>, - authentications: FuturesUnordered>>, + command_refresh: Option>>, + advertised_servers: Option>, idle_commands: JoinSet<()>, hooks: SessionHooks, idle_deadline: Option, } -/// List a runtime's MCP prompts as ACP available commands, de-duplicated by -/// name. Used at actor startup and after agent switches. -async fn available_commands_for(mcp: McpHandle) -> Result, SessionError> { - let prompts = mcp.list_prompts().await.map_err(SessionError::McpOperation)?; - let prompt_commands = prompts.iter().map(map_mcp_prompt_to_available_command).collect(); - Ok(dedupe_commands_by_name(prompt_commands)) -} - impl SessionActor { pub(crate) async fn spawn(init: SessionActorInit) -> Result { let cancel = CancellationToken::new(); let io = SessionIo::new(init.connection, init.session_id.clone()); let git_diff = GitDiffService::new(init.cwd.clone(), io.clone()); + let (elicitation_tx, elicitation_rx) = mpsc::channel(ELICITATION_CHANNEL_CAPACITY); let mut actor = SessionActor { io, git_diff, @@ -236,12 +231,14 @@ impl SessionActor { specs: init.specs, runtimes: HashMap::new(), runtime_factory: init.runtime_factory, + elicitation_tx, + elicitation_rx, transcript: init.transcript, config: init.config, modes: init.modes, prompts: PromptQueue::default(), command_refresh: None, - authentications: FuturesUnordered::new(), + advertised_servers: None, idle_commands: JoinSet::new(), hooks: init.hooks, idle_deadline: None, @@ -258,13 +255,11 @@ impl SessionActor { } async fn run(mut self, mut cmd_rx: mpsc::Receiver) { - if let Ok(runtime) = self.active_runtime() { - send_mcp_server_status(&self.io, runtime.mcp_server_statuses()); - } - self.refresh_available_commands(); + self.publish_active_mcps(); let shutdown = self.cancel.clone(); loop { let runtime = self.runtimes.get_mut(&self.active_agent).expect("active runtime is running"); + let (agent_events, mcp_changes) = runtime.events(); tokio::select! { biased; () = shutdown.cancelled() => { @@ -288,27 +283,18 @@ impl SessionActor { Some((message_id, content)) = self.prompts.next_prepared(), if self.prompts.is_preparing() => { self.submit_prompt(message_id, content).await; } - Some(result) = self.authentications.next(), if !self.authentications.is_empty() => { - if let Err(error) = result { - error!("MCP server authentication failed: {error}"); - } - } - result = async { self.command_refresh.as_mut().unwrap().await }, if self.command_refresh.is_some() => { + commands = async { self.command_refresh.as_mut().unwrap().await }, if self.command_refresh.is_some() => { self.command_refresh = None; - match result { - Ok(commands) => send_available_commands(&self.io, commands), - Err(error) => error!("Failed to refresh available commands: {error}"), - } + let update = acp::AvailableCommandsUpdate::new(commands); + self.io.send_update(SessionUpdate::AvailableCommandsUpdate(update)); } - Some(event) = runtime.agent_rx.recv() => { + Some(event) = agent_events.recv() => { self.on_agent_event(event).await; } - Some(event) = runtime.event_rx.recv() => { - let refresh_commands = matches!(event, McpClientEvent::ConnectionReady(_)); - on_mcp_client_event(&self.io, event); - if refresh_commands { - self.refresh_available_commands(); - } + Some(elicitation) = self.elicitation_rx.recv() => on_elicitation(&self.io, elicitation), + Some(()) = mcp_changes.next() => { + let statuses = self.active_runtime().mcp().statuses(); + self.publish_mcps(statuses); } else => break, } @@ -319,14 +305,22 @@ impl SessionActor { } } - fn refresh_available_commands(&mut self) { - if let Ok(runtime) = self.active_runtime() { - self.command_refresh = Some(available_commands_for(runtime.mcp().clone()).boxed()); + fn refresh_available_commands(&mut self, servers: &[McpServerStatusEntry]) { + let connected = servers + .iter() + .filter(|server| matches!(server.status, McpServerStatus::Connected { .. })) + .map(|server| server.name.clone()) + .collect::>(); + if self.advertised_servers.as_ref() == Some(&connected) { + return; } + let mcp = self.active_runtime().mcp().clone(); + self.command_refresh = Some(async move { available_commands(&prompt_commands(&mcp).await) }.boxed()); + self.advertised_servers = Some(connected); } - fn active_runtime(&self) -> Result<&AgentRuntime, SessionError> { - self.runtimes.get(&self.active_agent).ok_or(SessionError::ActiveRuntimeNotRunning) + fn active_runtime(&self) -> &AgentRuntime { + self.runtimes.get(&self.active_agent).expect("the active agent's runtime runs for the whole session") } fn active_provider_connections(&self) -> ProviderConnectionOverrides { @@ -352,11 +346,11 @@ impl SessionActor { async fn sync_active_conversation(&mut self) -> Result<(), SessionError> { let messages = conversation_messages_from_events(&self.transcript); - self.active_runtime()?.replace_conversation(messages).await + self.active_runtime().replace_conversation(messages).await } async fn send_active_command(&mut self, command: Command) -> Result<(), SessionError> { - self.active_runtime()?.send_agent_command(command).await + self.active_runtime().send_agent_command(command).await } async fn ensure_active_running(&mut self) -> Result<(), SessionError> { @@ -378,7 +372,7 @@ impl SessionActor { let runtime = tokio::select! { biased; () = self.cancel.cancelled() => return Err(SessionError::Cancelled), - runtime = self.runtime_factory.spawn(target.clone(), spec, messages, usage_seed) => runtime?, + runtime = self.runtime_factory.spawn(target.clone(), spec, messages, usage_seed, self.elicitation_tx.clone()) => runtime?, }; self.runtimes.insert(target.clone(), runtime); Ok(()) @@ -425,7 +419,7 @@ impl SessionActor { StateUpdate::Idle(IdleStateUpdate::new()) }; self.io.send_update(acp::SessionUpdate::StateUpdate(state)); - let _ = self.publish_active_mcps(); + self.publish_active_mcps(); let options = self.config.config_options(&self.modes, &available, self.oauth_credential_store.as_ref()); let _ = reply.send(Ok(options)); } @@ -448,13 +442,12 @@ impl SessionActor { self.io.send_update(acp::SessionUpdate::ConfigOptionUpdate(acp::ConfigOptionUpdate::new(options))); } SessionCommand::AuthenticateMcp { server_name } => { - if let Ok(runtime) = self.active_runtime() { - let mcp = runtime.mcp().clone(); - self.authentications.push( - async move { mcp.authenticate_server(&server_name).await.map_err(SessionError::McpOperation) } - .boxed(), - ); - } + let mcp = self.active_runtime().mcp().clone(); + tokio::spawn(async move { + if let Err(error) = mcp.authenticate(&server_name).await { + warn!(%error, server = server_name, "MCP server authentication failed"); + } + }); } SessionCommand::GitDiff { command } => self.git_diff.command(command).await, SessionCommand::GitDiffClose => self.git_diff.close().await, @@ -569,14 +562,14 @@ impl SessionActor { } } - async fn prepare_prompt_runtime(&mut self, idle: bool) -> Result { + async fn prepare_prompt_runtime(&mut self, idle: bool) -> Result { if idle { let switch = self.config.begin_prompt(&self.modes); self.apply_switch(switch).await?; self.send_active_command(Command::agent(AgentCommand::SetReasoningEffort(self.config.reasoning_effort))) .await?; } - Ok(self.active_runtime()?.mcp().clone()) + Ok(self.active_runtime().mcp().clone()) } async fn apply_deferred_agent_switch(&mut self) -> Result<(), SessionError> { @@ -612,7 +605,8 @@ impl SessionActor { if let Some(event) = self.select_agent(&agent_name).await? { self.persist_event(event); } - self.publish_active_mcps() + self.publish_active_mcps(); + Ok(()) } Switch::Model(model) => { let parser = ModelProviderParser::default() @@ -625,10 +619,15 @@ impl SessionActor { } } - fn publish_active_mcps(&mut self) -> Result<(), SessionError> { - send_mcp_server_status(&self.io, self.active_runtime()?.mcp_server_statuses()); - self.refresh_available_commands(); - Ok(()) + fn publish_active_mcps(&mut self) { + let statuses = self.active_runtime().mcp().statuses(); + self.advertised_servers = None; + self.publish_mcps(statuses); + } + + fn publish_mcps(&mut self, servers: Vec) { + self.refresh_available_commands(&servers); + self.io.send(McpNotification::ServerStatus { servers }); } fn record_agent_event(&mut self, message: &AgentEvent) { @@ -678,30 +677,15 @@ fn stop_reason(outcome: &TurnOutcome) -> acp::StopReason { } } -fn send_mcp_server_status(io: &SessionIo, servers: Vec) { - io.send(McpNotification::ServerStatus { servers }); -} - -fn on_mcp_client_event(io: &SessionIo, event: McpClientEvent) { - match event { - McpClientEvent::Elicitation(elicitation) => { +fn on_elicitation(io: &SessionIo, elicitation: Elicitation) { + match elicitation { + Elicitation::Request(elicitation) => { if let Some(connection) = &io.connection { - spawn_elicitation_request(connection, &io.session_id, *elicitation); - } else { - let _ = elicitation.response_sender.send(cancel_result()); + spawn_elicitation_request(connection, &io.session_id, elicitation); } } - McpClientEvent::ElicitationComplete { server_name, elicitation_id } => { - io.send(elicitation::build_acp_elicitation_completion_notification( - &io.session_id, - &server_name, - &elicitation_id, - )); - } - McpClientEvent::ServerStatusesChanged(servers) => send_mcp_server_status(io, servers), - McpClientEvent::ConnectionReady(snapshot) => send_mcp_server_status(io, snapshot.server_statuses()), - McpClientEvent::AuthenticationFailed { server, error } => { - error!("MCP server authentication failed for '{server}': {error}"); + Elicitation::Complete { server, id } => { + io.send(elicitation::build_acp_elicitation_completion_notification(&io.session_id, &server, &id)); } } } @@ -709,30 +693,27 @@ fn on_mcp_client_event(io: &SessionIo, event: McpClientEvent) { async fn on_elicitation_request( connection: &ConnectionTo, session_id: &SessionId, - elicitation: ElicitationRequest, + elicitation: Box, ) { let result = async { let request = - elicitation::map_mcp_elicitation_request_to_acp(&elicitation.server_name, session_id, &elicitation.request) + elicitation::map_mcp_elicitation_request_to_acp(&elicitation.server, session_id, &elicitation.request) .map_err(|error| error.to_string())?; let response = connection.send_request(request).block_task().await.map_err(|error| format!("{error:?}"))?; elicitation::map_acp_elicitation_response_to_mcp(response).map_err(|error| error.to_string()) } - .await - .unwrap_or_else(|error| { - error!("ACP elicitation failed: {error}"); - cancel_result() - }); + .await; - if elicitation.response_sender.send(result).is_err() { - error!("Failed to send elicitation response: receiver dropped"); + match result { + Ok(result) => elicitation.respond(result), + Err(error) => error!("ACP elicitation failed: {error}"), } } fn spawn_elicitation_request( connection: &ConnectionTo, session_id: &SessionId, - elicitation: ElicitationRequest, + elicitation: Box, ) { let connection = connection.clone(); let session_id = session_id.clone(); @@ -907,174 +888,4 @@ mod tests { assert_eq!(s.selected_mode.as_deref(), Some("Planner")); assert_eq!(s.effective_model(&validated_modes()), DEEPSEEK); } - - mod connection_tests { - use super::*; - use acp_utils::elicitation::source_mcp_server_name; - use acp_utils::testing::test_connection; - use rmcp::model::ElicitRequestParams; - use tokio::sync::oneshot; - use tokio::task::LocalSet; - use utils::mcp_status::McpServerStatus; - - fn dispatch_event(connection: &ConnectionTo, event: McpClientEvent) { - on_mcp_client_event(&SessionIo::new(Some(connection.clone()), SessionId::new("session-1")), event); - } - - #[tokio::test(flavor = "current_thread")] - async fn server_status_change_forwards_status_notification() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - let servers = - vec![McpServerStatusEntry::new("github", McpServerStatus::Connected { tool_count: 1 })]; - - dispatch_event(&cx, McpClientEvent::ServerStatusesChanged(servers)); - - let received = peer.next_mcp_notification().await; - assert!(matches!(received, McpNotification::ServerStatus { .. })); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn auth_failure_after_status_change_still_forwards_status() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - let servers = vec![McpServerStatusEntry::new( - "github", - McpServerStatus::Failed { error: "authentication timed out after 3 minutes".to_string() }, - )]; - - dispatch_event(&cx, McpClientEvent::ServerStatusesChanged(servers)); - dispatch_event( - &cx, - McpClientEvent::AuthenticationFailed { - server: "github".to_string(), - error: "authentication timed out after 3 minutes".to_string(), - }, - ); - - assert!(matches!(peer.next_mcp_notification().await, McpNotification::ServerStatus { .. })); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn empty_server_status_change_forwards_clear_notification() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - - dispatch_event(&cx, McpClientEvent::ServerStatusesChanged(vec![])); - - let McpNotification::ServerStatus { servers } = peer.next_mcp_notification().await; - assert_eq!(servers.len(), 0); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn status_event_forwards_server_status() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - let servers = - vec![McpServerStatusEntry::new("github", McpServerStatus::Connected { tool_count: 1 })]; - dispatch_event(&cx, McpClientEvent::ServerStatusesChanged(servers)); - - let McpNotification::ServerStatus { servers } = peer.next_mcp_notification().await; - assert_eq!(servers[0].name, "github"); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn elicitation_completion_forwards_native_acp_notification() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - - dispatch_event( - &cx, - McpClientEvent::ElicitationComplete { - server_name: "github".to_string(), - elicitation_id: "el-1".to_string(), - }, - ); - - let completion = peer.next_elicitation_completion().await; - assert_eq!(&*completion.elicitation_id.0, r#"["session-1","github","el-1"]"#); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn elicitation_request_forwards_response_from_peer() { - LocalSet::new() - .run_until(async { - let (cx, mut peer) = test_connection().await; - peer.queue_elicitation_response( - serde_json::from_value(serde_json::json!({ - "action": "accept", - "content": { "color": "red" } - })) - .unwrap(), - ); - - let (tx, rx) = oneshot::channel(); - let elicitation = ElicitationRequest { - server_name: "test-server".to_string(), - request: ElicitRequestParams::FormElicitationParams { - meta: None, - message: "Pick a color".to_string(), - requested_schema: rmcp::model::ElicitationSchema::builder() - .required_bool("approved") - .build() - .unwrap(), - }, - response_sender: tx, - }; - - on_elicitation_request(&cx, &SessionId::new("session-1"), elicitation).await; - - let result = rx.await.expect("response forwarded"); - assert_eq!(result.action, rmcp::model::ElicitationAction::Accept); - assert_eq!(result.content, Some(serde_json::json!({ "color": "red" }))); - - let received = peer.next_elicitation_request().await; - assert_eq!(source_mcp_server_name(received.meta.as_ref()), Some("test-server")); - let acp::ElicitationMode::Form(form) = received.mode else { panic!("expected form") }; - let acp::ElicitationScope::Session(scope) = form.scope else { panic!("expected session scope") }; - assert_eq!(&*scope.session_id.0, "session-1"); - }) - .await; - } - - #[tokio::test(flavor = "current_thread")] - async fn elicitation_request_surfaces_cancel_on_transport_error() { - LocalSet::new() - .run_until(async { - let (cx, _peer) = test_connection().await; - let (tx, rx) = oneshot::channel(); - let elicitation = ElicitationRequest { - server_name: "test-server".to_string(), - request: ElicitRequestParams::UrlElicitationParams { - meta: None, - message: "Authorize".to_string(), - url: "https://example.com".to_string(), - elicitation_id: "el-1".to_string(), - }, - response_sender: tx, - }; - - on_elicitation_request(&cx, &SessionId::new("session-1"), elicitation).await; - - let result = rx.await.expect("response forwarded"); - assert_eq!(result.action, rmcp::model::ElicitationAction::Cancel); - }) - .await; - } - } } diff --git a/crates/aether-cli/src/acp/session/error.rs b/crates/aether-cli/src/acp/session/error.rs index 4c2874f89..086ccec84 100644 --- a/crates/aether-cli/src/acp/session/error.rs +++ b/crates/aether-cli/src/acp/session/error.rs @@ -1,26 +1,19 @@ use thiserror::Error; -use crate::error::CliError; +use crate::runtime::RuntimeError; #[derive(Debug, Error)] pub enum SessionError { #[error("failed to build session: {0}")] - Build(#[from] CliError), + Build(#[from] RuntimeError), #[error("failed to persist session: {0}")] Persistence(#[from] aether_sessions::SessionStoreError), #[error("agent command channel is closed")] CommandChannelClosed, - #[error("MCP operation failed: {0}")] - McpOperation(#[from] aether_core::mcp::McpHandleError), #[error("model configuration failed: {0}")] Model(#[from] llm::LlmError), - #[cfg(any(test, feature = "testing"))] - #[error("MCP runtime stopped during startup")] - McpStartupStopped, #[error("agent runtime not found: {0}")] AgentNotFound(String), #[error("session shut down during runtime startup")] Cancelled, - #[error("active agent runtime is not running")] - ActiveRuntimeNotRunning, } diff --git a/crates/aether-cli/src/acp/session/factory.rs b/crates/aether-cli/src/acp/session/factory.rs index ef78e92d8..773965d9d 100644 --- a/crates/aether-cli/src/acp/session/factory.rs +++ b/crates/aether-cli/src/acp/session/factory.rs @@ -9,9 +9,9 @@ use agent_client_protocol::{Client, ConnectionTo, Error}; use llm::catalog::{LlmModel, get_local_models}; use llm::types::IsoString; use llm::{ProviderConnectionOverrides, ReasoningEffort}; -use rmcp::model::ClientCapabilities; +use mcp_utils::model::ElicitationCapability; use std::path::{Path, PathBuf}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use tokio::sync::OnceCell; use tracing::{error, info, warn}; @@ -39,6 +39,7 @@ pub(crate) struct SessionFactory { runtime_factory: Option>, hooks: SessionHooks, available: OnceCell>, + elicitation: OnceLock, } /// The fully-built session ready to be registered with [`AcpState`](crate::acp::state::AcpState). @@ -91,9 +92,14 @@ impl SessionFactory { runtime_factory, hooks, available: OnceCell::new(), + elicitation: OnceLock::new(), } } + pub(crate) fn set_elicitation(&self, capability: ElicitationCapability) { + let _ = self.elicitation.set(capability); + } + pub(crate) async fn available_models(&self) -> &[LlmModel] { self.available.get_or_init(get_local_models).await } @@ -102,7 +108,6 @@ impl SessionFactory { &self, mut args: NewSessionRequest, cx: Option<&ConnectionTo>, - mcp_capabilities: ClientCapabilities, ) -> Result { let cwd = args.cwd.clone().into_inner(); let mcp_servers = args.mcp_servers.clone(); @@ -140,7 +145,6 @@ impl SessionFactory { args.cwd.into_inner(), args.mcp_servers, mode_catalog.specs.catalog(), - mcp_capabilities, &session_id, ) }); @@ -161,10 +165,9 @@ impl SessionFactory { &self, args: ResumeSessionRequest, cx: &ConnectionTo, - mcp_capabilities: ClientCapabilities, replay: bool, ) -> Result { - self.restore(args.session_id, args.cwd.into_inner(), args.mcp_servers, cx, mcp_capabilities, replay).await + self.restore(args.session_id, args.cwd.into_inner(), args.mcp_servers, cx, replay).await } async fn restore( @@ -173,7 +176,6 @@ impl SessionFactory { cwd: PathBuf, mcp_servers: Vec, cx: &ConnectionTo, - mcp_capabilities: ClientCapabilities, replay: bool, ) -> Result { let session_id_string = session_id.0.to_string(); @@ -197,7 +199,6 @@ impl SessionFactory { cwd.clone(), mcp_servers.clone(), mode_catalog.specs.catalog(), - mcp_capabilities, session_id.0.as_ref(), ) }); @@ -219,14 +220,18 @@ impl SessionFactory { cwd: PathBuf, mcp_servers: Vec, catalog: &AgentCatalog, - mcp_capabilities: ClientCapabilities, session_affinity_key: &str, ) -> Arc { let deps = AgentDeps::new(Arc::clone(&self.oauth_credential_store), self.observer_factory.clone()) .with_agent_registry(catalog.registry().clone()) - .with_mcp_client_capabilities(mcp_capabilities) .with_session_affinity_key(session_affinity_key); - Arc::new(ProductionRuntimeFactory::new(cwd, map_acp_mcp_servers(mcp_servers), deps)) + let mcp_servers = map_acp_mcp_servers(mcp_servers, &cwd); + Arc::new(ProductionRuntimeFactory::new( + cwd, + mcp_servers, + deps, + self.elicitation.get().cloned().unwrap_or_default(), + )) } #[allow(clippy::too_many_arguments)] diff --git a/crates/aether-cli/src/acp/session/mod.rs b/crates/aether-cli/src/acp/session/mod.rs index 33b5cb0d3..d0eb93ef1 100644 --- a/crates/aether-cli/src/acp/session/mod.rs +++ b/crates/aether-cli/src/acp/session/mod.rs @@ -12,7 +12,6 @@ mod pagination; pub(crate) mod prompts; pub(crate) mod registry; pub(crate) mod runtime; -pub(crate) mod slash_commands; pub(crate) use pagination::paginate_summaries; pub(crate) use registry::SessionRegistry; diff --git a/crates/aether-cli/src/acp/session/model.rs b/crates/aether-cli/src/acp/session/model.rs index 4f39d29e6..d70b2c4e6 100644 --- a/crates/aether-cli/src/acp/session/model.rs +++ b/crates/aether-cli/src/acp/session/model.rs @@ -253,9 +253,9 @@ mod tests { use super::*; use aether_auth::FakeOAuthCredentialStore; use aether_core::agent_spec::AgentSpecExposure; + use aether_core::mcp::{McpConfig, ToolFilter}; use agent_client_protocol::schema::v2::{SessionConfigKind, SessionConfigSelectOption, SessionConfigSelectOptions}; use llm::catalog::{AnthropicModel, BedrockFoundationModel, BedrockModel, DeepSeekModel, GeminiModel}; - use mcp_utils::client::ToolFilter; fn test_models() -> Vec { vec![ @@ -277,7 +277,7 @@ mod tests { context_window: None, prompts: vec![], provider_connections: llm::ProviderConnectionOverrides::default(), - mcp_config_sources: Vec::new(), + mcp: McpConfig::default(), exposure: AgentSpecExposure::both(), tools: ToolFilter::default(), } diff --git a/crates/aether-cli/src/acp/session/runtime.rs b/crates/aether-cli/src/acp/session/runtime.rs index 9c1f008a0..8dd5bac80 100644 --- a/crates/aether-cli/src/acp/session/runtime.rs +++ b/crates/aether-cli/src/acp/session/runtime.rs @@ -1,74 +1,50 @@ use super::agent_key::AgentKey; use super::error::SessionError; use crate::runtime::{Runtime, RuntimeBuilder}; -use aether_auth::OAuthHandler; use aether_core::agent_spec::AgentSpec; -use aether_core::core::{AgentDeps, AgentHandle}; +use aether_core::core::AgentDeps; use aether_core::events::{AgentCommand, AgentEvent, Command}; -use aether_core::mcp::{McpHandle, McpRuntime}; +use aether_core::mcp::McpCatalog; +use futures::Stream; use llm::{ChatMessage, SessionUsageEvent}; -use mcp_utils::client::{ - ElicitingOAuthHandler, McpClientEvent, McpConnectionDetails, McpError, McpServer, OAuthHandlerFactory, -}; +use mcp_utils::client::{Elicitation, ElicitationSink}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::ElicitationCapability; use std::path::PathBuf; -use std::sync::Arc; -use tokio::sync::{mpsc, watch}; -use utils::mcp_status::McpServerStatusEntry; +use std::pin::Pin; +use tokio::sync::mpsc; pub(crate) struct AgentRuntime { - pub(crate) agent_rx: mpsc::Receiver, - pub(crate) event_rx: mpsc::Receiver, - agent_tx: mpsc::Sender, - latest_mcp_snapshot: watch::Receiver, - agent_handle: Option, - mcp_runtime: McpRuntime, + runtime: Runtime, + mcp_changes: McpChanges, } +pub(crate) type McpChanges = Pin + Send + Sync>>; + impl AgentRuntime { - pub(crate) fn new( - agent_tx: mpsc::Sender, - agent_rx: mpsc::Receiver, - agent_handle: Option, - event_rx: mpsc::Receiver, - mcp_runtime: McpRuntime, - ) -> Self { - let latest_mcp_snapshot = mcp_runtime.handle().subscribe(); - Self { agent_rx, event_rx, agent_tx, latest_mcp_snapshot, agent_handle, mcp_runtime } + pub(crate) fn new(runtime: Runtime) -> Self { + let mcp_changes = Box::pin(runtime.mcp.changes()); + Self { runtime, mcp_changes } } - pub(crate) async fn shutdown(mut self) { - if let Some(handle) = self.agent_handle.take() { - handle.abort(); - handle.await_completion().await; - } - self.mcp_runtime.shutdown().await; + pub(crate) fn events(&mut self) -> (&mut mpsc::Receiver, &mut McpChanges) { + (&mut self.runtime.agent_rx, &mut self.mcp_changes) } - pub(crate) async fn send_agent_command(&self, command: Command) -> Result<(), SessionError> { - self.agent_tx.send(command).await.map_err(|_| SessionError::CommandChannelClosed) + pub(crate) async fn shutdown(self) { + self.runtime.shutdown().await; } - pub(crate) async fn replace_conversation(&self, messages: Vec) -> Result<(), SessionError> { - self.agent_tx - .send(Command::agent(AgentCommand::ReplaceConversation(messages))) - .await - .map_err(|_| SessionError::CommandChannelClosed) + pub(crate) async fn send_agent_command(&self, command: Command) -> Result<(), SessionError> { + self.runtime.agent_tx.send(command).await.map_err(|_| SessionError::CommandChannelClosed) } - pub(crate) fn mcp(&self) -> &McpHandle { - self.mcp_runtime.handle() + pub(crate) async fn replace_conversation(&self, messages: Vec) -> Result<(), SessionError> { + self.send_agent_command(Command::agent(AgentCommand::ReplaceConversation(messages))).await } - pub(crate) fn mcp_server_statuses(&self) -> Vec { - self.latest_mcp_snapshot.borrow().server_statuses() - } -} - -impl Drop for AgentRuntime { - fn drop(&mut self) { - if let Some(handle) = &self.agent_handle { - handle.abort(); - } + pub(crate) fn mcp(&self) -> &McpCatalog { + &self.runtime.mcp } } @@ -83,18 +59,25 @@ pub(crate) trait RuntimeFactory: Send + Sync { spec: &AgentSpec, initial_messages: Vec, usage_seed: Option, + elicitations: mpsc::Sender, ) -> Result; } pub(crate) struct ProductionRuntimeFactory { cwd: PathBuf, - mcp_servers: Vec, + mcp_servers: Vec, agent_deps: AgentDeps, + elicitation: ElicitationCapability, } impl ProductionRuntimeFactory { - pub fn new(cwd: PathBuf, client_servers: Vec, agent_deps: AgentDeps) -> Self { - Self { cwd, mcp_servers: client_servers, agent_deps } + pub fn new( + cwd: PathBuf, + mcp_servers: Vec, + agent_deps: AgentDeps, + elicitation: ElicitationCapability, + ) -> Self { + Self { cwd, mcp_servers, agent_deps, elicitation } } } @@ -106,30 +89,16 @@ impl RuntimeFactory for ProductionRuntimeFactory { spec: &AgentSpec, initial_messages: Vec, usage_seed: Option, + elicitations: mpsc::Sender, ) -> Result { - let extra_servers = self.mcp_servers.clone(); - - let mut builder = RuntimeBuilder::from_spec(self.cwd.clone(), spec.clone()) - .extra_servers(extra_servers) - .agent_deps(self.agent_deps.clone()); - if let Some(last) = usage_seed { - builder = builder.resume_usage(last); - } - if self.agent_deps.supports_mcp_url_elicitation() { - builder = builder.oauth_handler_factory(mcp_oauth_handler_factory()); - } - - let runtime = builder.build(None, Some(initial_messages)).await?; - - let Runtime { agent_tx, agent_rx, agent_handle, event_rx, mcp_runtime } = runtime; - Ok(AgentRuntime::new(agent_tx, agent_rx, Some(agent_handle), event_rx, mcp_runtime)) + let sink = ElicitationSink::new(elicitations, self.elicitation.clone()); + let runtime = RuntimeBuilder::from_spec(self.cwd.clone(), spec.clone()) + .extra_servers(self.mcp_servers.clone()) + .elicitation(sink) + .agent_deps(self.agent_deps.clone()) + .resume_usage(usage_seed) + .build(initial_messages) + .await?; + Ok(AgentRuntime::new(runtime)) } } - -fn mcp_oauth_handler_factory() -> OAuthHandlerFactory { - Arc::new(|ctx| { - ElicitingOAuthHandler::new(ctx) - .map(|handler| Arc::new(handler) as Arc) - .map_err(|error| McpError::ConnectionFailed(format!("failed to initialize OAuth handler: {error}"))) - }) -} diff --git a/crates/aether-cli/src/acp/session/slash_commands.rs b/crates/aether-cli/src/acp/session/slash_commands.rs deleted file mode 100644 index 225503b65..000000000 --- a/crates/aether-cli/src/acp/session/slash_commands.rs +++ /dev/null @@ -1,39 +0,0 @@ -use super::actor::SessionIo; -use agent_client_protocol::schema::v2 as acp; -use llm::ContentBlock; -use tracing::{error, info}; - -use crate::slash_commands::{expand_slash_command, parse_slash_command}; -use aether_core::mcp::McpHandle; - -pub(crate) async fn expand_slash_command_in_content( - mcp: &McpHandle, - mut content: Vec, -) -> Vec { - if let Some(ContentBlock::Text { text }) = content.first() { - let expanded = expand_slash_command_text(mcp, text.clone()).await; - content[0] = ContentBlock::text(expanded); - } - content -} - -async fn expand_slash_command_text(mcp: &McpHandle, text: String) -> String { - let Some(slash_command) = parse_slash_command(&text) else { - return text; - }; - - match expand_slash_command(mcp, slash_command.command_name, slash_command.args_text).await { - Ok(expanded) => { - info!("Expanded slash command -> {} chars", expanded.len()); - expanded - } - Err(e) => { - error!("Failed to expand slash command: {}", e); - text - } - } -} - -pub(crate) fn send_available_commands(io: &SessionIo, available_commands: Vec) { - io.send_update(acp::SessionUpdate::AvailableCommandsUpdate(acp::AvailableCommandsUpdate::new(available_commands))); -} diff --git a/crates/aether-cli/src/acp/state.rs b/crates/aether-cli/src/acp/state.rs index 9fd15af9d..3eafcb77f 100644 --- a/crates/aether-cli/src/acp/state.rs +++ b/crates/aether-cli/src/acp/state.rs @@ -21,12 +21,11 @@ use agent_client_protocol::util::internal_error; use agent_client_protocol::{Agent, Client, ConnectTo, ConnectionTo, Error, Responder}; use llm::catalog::{LlmModel, ModelSpec}; use llm::{ContentBlock, ProviderConnectionOverrides}; -use mcp_utils::client::{client_capabilities, client_capabilities_for}; -use rmcp::model::ClientCapabilities; +use mcp_utils::model::{ElicitationCapability, FormElicitationCapability, UrlElicitationCapability}; use std::collections::HashSet; use std::path::PathBuf; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, OnceLock}; use tokio::sync::{mpsc, oneshot, watch}; use tokio::task::spawn_blocking; use tokio_util::sync::CancellationToken; @@ -62,7 +61,6 @@ pub(crate) struct AcpState { oauth_credential_store: Arc, factory: SessionFactory, telemetry: Option>, - mcp_capabilities: OnceLock, cwd: PathBuf, } @@ -187,14 +185,13 @@ impl AcpState { oauth_credential_store: config.oauth_credential_store, factory, telemetry: config.telemetry, - mcp_capabilities: OnceLock::new(), cwd: config.cwd, } } pub(crate) async fn initialize(&self, args: InitializeRequest) -> Result { info!("Received initialize request: {:?}", args); - let _ = self.mcp_capabilities.set(mcp_client_capabilities(&args.capabilities)); + self.factory.set_elicitation(mcp_elicitation(&args.capabilities)); let auth_methods = build_auth_methods(self.oauth_credential_store.as_ref()); let available = self.factory.available_models().await.to_vec(); let prompt_capabilities = prompt_capabilities_for_models(&available); @@ -267,8 +264,7 @@ impl AcpState { request: NewSessionRequest, connection: Option<&ConnectionTo>, ) -> Result { - let mcp_capabilities = self.mcp_capabilities(); - let prepared = self.factory.prepare_new(request, connection, mcp_capabilities).await?; + let prepared = self.factory.prepare_new(request, connection).await?; self.registry.stop().await; let created = prepared.start().await?; let spawned = SpawnedSession { @@ -321,8 +317,7 @@ impl AcpState { return Ok(ResumeSessionResponse::new().config_options(options)); } - let mcp_capabilities = self.mcp_capabilities(); - let prepared = self.factory.prepare_resume(req, cx, mcp_capabilities, replay).await?; + let prepared = self.factory.prepare_resume(req, cx, replay).await?; self.registry.stop().await; let created = prepared.start().await?; let response = ResumeSessionResponse::new().config_options(created.config_options); @@ -610,10 +605,6 @@ impl AcpState { self.broadcast_config_options().await; } - fn mcp_capabilities(&self) -> ClientCapabilities { - self.mcp_capabilities.get().cloned().unwrap_or_else(client_capabilities) - } - async fn broadcast_config_options(&self) { if let Some(sender) = self.registry.lookup(None).await { let available = self.factory.available_models().await.to_vec(); @@ -657,12 +648,16 @@ fn workspace_error(e: &WorkspaceError) -> Error { if e.is_invalid_input() { Error::invalid_params().data(e.to_string()) } else { internal_error(e.to_string()) } } -fn mcp_client_capabilities(client: &acp::ClientCapabilities) -> rmcp::model::ClientCapabilities { +fn mcp_elicitation(client: &acp::ClientCapabilities) -> ElicitationCapability { let elicitation = client.elicitation.as_ref(); - client_capabilities_for( - elicitation.is_some_and(|capabilities| capabilities.form.is_some()), - elicitation.is_some_and(|capabilities| capabilities.url.is_some()), - ) + let mut capability = ElicitationCapability::new(); + if elicitation.is_some_and(|capabilities| capabilities.form.is_some()) { + capability = capability.with_form(FormElicitationCapability::new()); + } + if elicitation.is_some_and(|capabilities| capabilities.url.is_some()) { + capability = capability.with_url(UrlElicitationCapability::new()); + } + capability } fn build_auth_methods(store: &dyn OAuthCredentialStorage) -> Vec { @@ -811,20 +806,19 @@ mod tests { #[test] fn mcp_elicitation_capabilities_mirror_the_acp_client() { - let none = mcp_client_capabilities(&acp::ClientCapabilities::new()); - assert!(none.elicitation.is_none()); + let none = mcp_elicitation(&acp::ClientCapabilities::new()); + assert_eq!(none, ElicitationCapability::new()); let form_only = acp::ClientCapabilities::new() .elicitation(acp::ElicitationCapabilities::new().form(acp::ElicitationFormCapabilities::new())); - let form = mcp_client_capabilities(&form_only).elicitation.unwrap(); - assert!(form.form.is_some()); - assert!(form.url.is_none()); + assert_eq!( + mcp_elicitation(&form_only), + ElicitationCapability::new().with_form(FormElicitationCapability::new()) + ); let url_only = acp::ClientCapabilities::new() .elicitation(acp::ElicitationCapabilities::new().url(acp::ElicitationUrlCapabilities::new())); - let url = mcp_client_capabilities(&url_only).elicitation.unwrap(); - assert!(url.form.is_none()); - assert!(url.url.is_some()); + assert_eq!(mcp_elicitation(&url_only), ElicitationCapability::new().with_url(UrlElicitationCapability::new())); } #[test] diff --git a/crates/aether-cli/src/acp/testing.rs b/crates/aether-cli/src/acp/testing.rs index 6981d0094..f1aae9944 100644 --- a/crates/aether-cli/src/acp/testing.rs +++ b/crates/aether-cli/src/acp/testing.rs @@ -1,4 +1,3 @@ -use super::fake_prompt_mcp::FakePromptMcp; use super::server::{AcpServer, ServerRunError}; use super::session::actor::SessionActorInit; use super::session::agent_key::AgentKey; @@ -9,8 +8,8 @@ use super::session::hooks::SessionHooks; use super::session::model::{Modes, ValidatedMode}; use super::session::runtime::{AgentRuntime, RuntimeFactory}; use super::state::{AcpState, AcpStateConfig}; -use crate::error::CliError; use crate::resolve::InitialSessionSelection; +use crate::runtime::{Runtime, RuntimeBuilder}; use crate::settings_args::SettingsSourceArgs; use crate::workspace::WorkspaceManager; use crate::workspace::testing::StdCopyCloner; @@ -20,24 +19,23 @@ use aether_auth::OAuthCredentialStorage; use aether_core::agent_spec::{AgentSpec, AgentSpecExposure}; use aether_core::core::{AgentBuilder, AgentHandle, Prompt}; use aether_core::events::{AgentEvent, Command, MessageEvent, ToolEvent, TurnEvent, TurnOutcome, UserCommand}; -use aether_core::mcp::{McpBuilder, McpSession, ServerFactory, mcp}; +use aether_core::mcp::{McpBuilder, McpCatalog, mcp}; +use aether_core::testing::ready_catalog; use aether_project::AgentCatalog; use aether_sessions::SessionStore; use aether_sessions::{SessionControlEvent, SessionEvent, SessionMeta, UserEvent, last_agent_from_events}; use agent_client_protocol::schema::v2::{ - AbsolutePath, InitializeResponse, ReplayFrom, ReplayFromStart, ResumeSessionRequest, SessionId, SessionUpdate, - StateUpdate, StopReason, + AbsolutePath, AvailableCommand, InitializeResponse, ReplayFrom, ReplayFromStart, ResumeSessionRequest, SessionId, + SessionUpdate, StateUpdate, StopReason, }; use agent_client_protocol::{Agent, Channel, Client, ConnectionTo, on_receive_notification}; -use futures::FutureExt; use llm::testing::FakeLlmProvider; use llm::{ChatMessage, Context, LlmModel, LlmResponse, SessionUsageEvent, StreamingModelProvider}; use llm::{MessageId, ProviderConnectionOverrides, ToolCallRequest, ToolCallResult}; -use mcp_utils::client::{InMemoryServerSpec, McpServer, McpTransport, ToolExposure}; -use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse, completed_task_payload}; -use rmcp::RoleServer; -use rmcp::model::{CallToolResult, CreateTaskResult, DetailedTask, Task, TaskPayload, TaskStatus}; -use rmcp::service::DynService; +use mcp_utils::client::{Elicitation, ElicitationSink}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::{CallToolResult, CreateTaskResult, DetailedTask, ElicitResult, Task, TaskPayload, TaskStatus}; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse, Gate, completed_task_payload, elicitation_sink}; use std::collections::HashMap; use std::future::Future; use std::net::SocketAddr; @@ -46,6 +44,7 @@ use std::sync::{Arc, Mutex}; use tokio::sync::{mpsc, oneshot, watch}; use tokio::task::{JoinHandle, JoinSet, LocalSet}; use tokio_util::sync::CancellationToken; +use utils::mcp_status::McpServerStatusEntry; const PLANNER_REPLY: &str = "planner reply"; const CODER_REPLY: &str = "coder reply"; @@ -140,17 +139,25 @@ impl AcpTestHarness { PendingRuntime { started: observed, release } } - pub fn pause_prompt_expansion(&self) -> (Arc, Arc) { - let gate = (Arc::new(tokio::sync::Notify::new()), Arc::new(tokio::sync::Notify::new())); + pub fn pause_prompt_expansion(&self) -> Gate { + let gate = Gate::new(); self.runtime_control.lock().unwrap().prompt_gate = Some(gate.clone()); gate } + pub fn add_mcp_server(&self, server: ServerSpec) { + self.runtime_control.lock().unwrap().extra_servers.push(server); + } + + pub fn add_stalled_mcp_server(&self) { + self.runtime_control.lock().unwrap().stalled_server = true; + } + pub fn watch_agent_prompts(&self) -> watch::Receiver { self.runtime_control.lock().unwrap().agent_prompts.subscribe() } - pub fn elicit_during_prompt_expansion(&self) -> mpsc::UnboundedReceiver { + pub fn elicit_during_prompt_expansion(&self) -> mpsc::UnboundedReceiver { let (sender, receiver) = mpsc::unbounded_channel(); self.runtime_control.lock().unwrap().elicitation_results = Some(sender); receiver @@ -280,7 +287,11 @@ impl AcpTestHarness { workspace_manager, oauth_credential_store: oauth_store.clone(), initial_selection: InitialSessionSelection::default(), - settings_source: SettingsSourceArgs::default(), + // Inline settings keep the developer's own `~/.aether/settings.json` out of the tests. + settings_source: SettingsSourceArgs { + settings_json: Some(r#"{"agents":[]}"#.to_string()), + settings_file: None, + }, provider_connections: ProviderConnectionOverrides::default(), telemetry: None, runtime_factory: Some(runtime_factory), @@ -369,19 +380,39 @@ impl AcpTestHarness { } pub async fn expect_mcp_server_status(&mut self, expected: &[&str]) { - assert_server_status(self.peer.next_mcp_notification().await, expected); + let statuses = self.next_mcp_server_statuses().await; + let names = statuses.iter().map(|server| server.name.as_str()).collect::>(); + for server_name in expected { + assert!(names.contains(server_name), "expected server {server_name} in {names:?}"); + } } pub async fn expect_mcp_server_status_exact(&mut self, expected: &[&str]) { - assert_server_status_exact(self.peer.next_mcp_notification().await, expected); + let statuses = self.next_mcp_server_statuses().await; + let names = statuses.iter().map(|server| server.name.as_str()).collect::>(); + assert_eq!(names, expected); + } + + pub async fn next_mcp_server_statuses(&mut self) -> Vec { + let McpNotification::ServerStatus { servers } = self.peer.next_mcp_notification().await; + servers } pub async fn expect_available_commands(&mut self, expected: &[&str], unexpected: &[&str]) { + let commands = self.next_available_commands().await; + let names = commands.iter().map(|command| command.name.as_str()).collect::>(); + for name in expected { + assert!(names.contains(name), "expected command /{name} in {names:?}"); + } + for name in unexpected { + assert!(!names.contains(name), "did not expect command /{name} in {names:?}"); + } + } + + pub async fn next_available_commands(&mut self) -> Vec { loop { - let update = self.peer.next_session_notification().await.update; - if matches!(update, SessionUpdate::AvailableCommandsUpdate(_)) { - assert_available_commands(update, expected, unexpected); - return; + if let SessionUpdate::AvailableCommandsUpdate(update) = self.peer.next_session_notification().await.update { + return update.available_commands; } } } @@ -409,10 +440,9 @@ impl AcpTestHarness { id: SessionId, model: &str, ) { - let mut mcp = mcp("/tmp").spawn().await.expect("stub MCP spawns"); - mcp.block_until_ready().await.expect("stub MCP is ready"); - self.register_stub_session(StubAgentParts { tx: agent_tx, rx: agent_rx, handle: agent_handle, mcp }, id, model) - .await; + let (mcp, elicitations) = spawn_ready(mcp("/tmp")).await; + let runtime = Runtime { agent_tx, agent_rx, agent_handle, mcp }; + self.register_stub_session(runtime, elicitations, id, model).await; } pub async fn insert_background_task_session( @@ -420,22 +450,20 @@ impl AcpTestHarness { provider: FakeLlmProvider, task: &FakeBackgroundTask, id: SessionId, - ) { + ) -> McpCatalog { let server = task.server.clone(); - let gate = self.runtime_control.lock().unwrap().prompt_gate.clone(); - let mcp = with_in_memory_server(mcp("/tmp"), "tasks", move || server.clone().into_dyn()); - let mcp = with_in_memory_server(mcp, "prompts", move || { - FakePromptMcp::new("plan").with_gate(gate.clone()).into_dyn() - }); - let mut mcp = mcp.spawn().await.expect("background task MCP spawns"); - mcp.block_until_ready().await.expect("background task MCP is ready"); - let (tx, rx, handle) = AgentBuilder::new(Arc::new(provider)) + let prompts = self.runtime_control.lock().unwrap().prompt_server("plan"); + let servers = vec![ServerSpec::in_process("tasks", server), ServerSpec::in_process("prompts", prompts)]; + let (mcp, elicitations) = spawn_ready(mcp("/tmp").with_servers(servers)).await; + let (agent_tx, agent_rx, agent_handle) = AgentBuilder::new(Arc::new(provider)) .max_auto_continues(0) - .tools(mcp.handle().clone(), Vec::new()) + .mcp(mcp.clone()) .spawn() .await .expect("background task agent spawns"); - self.register_stub_session(StubAgentParts { tx, rx, handle, mcp }, id, "fake:fake").await; + let runtime = Runtime { agent_tx, agent_rx, agent_handle, mcp: mcp.clone() }; + self.register_stub_session(runtime, elicitations, id, "fake:fake").await; + mcp } pub fn append_stored_session(&self, session_id: &str, created_at: &str) { @@ -577,14 +605,20 @@ impl AcpTestHarness { self.session_store.append_event(session_id, event).expect("stored session event appends"); } - async fn register_stub_session(&self, mut parts: StubAgentParts, id: SessionId, model: &str) { + async fn register_stub_session( + &self, + mut runtime: Runtime, + elicitations: mpsc::Receiver, + id: SessionId, + model: &str, + ) { let model_spec: llm::catalog::LlmModel = "anthropic:claude-sonnet-4-5".parse().expect("test model parses"); let mut specs = SessionAgents::new(AgentCatalog::empty(PathBuf::from("/tmp"))); specs.set_default(AgentSpec::bare(&model_spec, None, Vec::new())); let agent_prompts = self.runtime_control.lock().unwrap().agent_prompts.clone(); - self.runtime_control.lock().unwrap().agents.push(parts.tx.clone()); - parts.tx = count_prompts(parts.tx, agent_prompts); - let factory = Arc::new(StubRuntimeFactory { agent_parts: Mutex::new(Some(parts)) }); + self.runtime_control.lock().unwrap().agents.push(runtime.agent_tx.clone()); + runtime.agent_tx = count_prompts(runtime.agent_tx, agent_prompts); + let factory = Arc::new(StubRuntimeFactory { runtime: Mutex::new(Some((runtime, elicitations))) }); self.state .register_session(SessionActorInit { @@ -671,18 +705,33 @@ impl FakeBackgroundTask { const ID: &str = "background-task"; pub fn new() -> Self { - let working = background_task(TaskStatus::Working); let server = FakeMcpServer::new() - .with_tool(FakeTool::new("start").responds(FakeToolResponse::task(CreateTaskResult::new(working.clone())))) - .with_task(Self::ID, [DetailedTask::new(working, TaskPayload::Working)]); + .with_tool(Self::start_tool()) + .with_task(Self::ID, [DetailedTask::new(background_task(TaskStatus::Working), TaskPayload::Working)]); Self { server } } pub fn complete(&self, output: &str) { - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text(output)]); + let result = CallToolResult::success(vec![mcp_utils::model::ContentBlock::text(output)]); let completed = DetailedTask::new(background_task(TaskStatus::Completed), completed_task_payload(result)); self.server.state().script_task(Self::ID, [completed]); } + + pub async fn announce_unchanged_tools(&self) { + let state = self.server.state(); + let listed = state.tool_list_requests(); + state.notify_tools_changed().await; + state.wait_until(|state| state.tool_list_requests() > listed).await; + } + + pub async fn add_tool(&self, name: &str) { + self.server.state().add_tool_and_notify(FakeTool::new(name)).await; + } + + fn start_tool() -> FakeTool { + let working = background_task(TaskStatus::Working); + FakeTool::new("start").responds(FakeToolResponse::task(CreateTaskResult::new(working))) + } } impl Default for FakeBackgroundTask { @@ -764,9 +813,37 @@ struct FakeRuntimeFactory { struct FakeRuntimeControl { agents: Vec>, pending: Option<(oneshot::Sender<()>, oneshot::Receiver)>, - prompt_gate: Option<(Arc, Arc)>, - elicitation_results: Option>, + prompt_gate: Option, + elicitation_results: Option>, agent_prompts: watch::Sender, + extra_servers: Vec, + stalled_server: bool, +} + +impl FakeRuntimeControl { + fn servers(&self, def: &FakeAgentDef) -> Vec { + let mut servers = def + .mcp + .iter() + .map(|(server_name, prompt_name)| ServerSpec::in_process(server_name, self.prompt_server(prompt_name))) + .collect::>(); + servers.extend(self.extra_servers.iter().cloned()); + if self.stalled_server { + servers.push(ServerSpec::in_process("stalled-mcp", FakeMcpServer::new().with_tool_list_gate(Gate::new()))); + } + servers + } + + fn prompt_server(&self, prompt_name: &str) -> FakeMcpServer { + let mut server = FakeMcpServer::new().with_prompt(prompt_name, format!("expanded {prompt_name}")); + if let Some(gate) = self.prompt_gate.clone() { + server = server.with_prompt_gate(gate); + } + if let Some(results) = self.elicitation_results.clone() { + server = server.with_prompt_elicitation(results); + } + server + } } struct FakeAgentDef { @@ -783,6 +860,7 @@ impl RuntimeFactory for FakeRuntimeFactory { spec: &AgentSpec, initial_messages: Vec, usage_seed: Option, + elicitations: mpsc::Sender, ) -> Result { let pending = self.control.lock().unwrap().pending.take(); if let Some((started, proceed)) = pending { @@ -796,53 +874,29 @@ impl RuntimeFactory for FakeRuntimeFactory { .get(&spec.name) .or_else(|| self.agents.values().next()) .ok_or_else(|| SessionError::AgentNotFound(spec.name.clone()))?; - let provider = def.provider.clone(); - - let mut mcp_builder = mcp(&self.cwd).with_tool_filter(spec.tools.clone()); - if let Some((server_name, prompt_name)) = &def.mcp { - let prompt_name = prompt_name.clone(); - let gate = self.control.lock().unwrap().prompt_gate.clone(); - let elicitation_results = self.control.lock().unwrap().elicitation_results.clone(); - mcp_builder = with_in_memory_server(mcp_builder, server_name, move || { - FakePromptMcp::new(&prompt_name) - .with_gate(gate.clone()) - .with_elicitation(elicitation_results.clone()) - .into_dyn() - }); - } - let mut spawn = - mcp_builder.spawn().await.map_err(|e| SessionError::Build(CliError::McpError(e.to_string())))?; - spawn.block_until_ready().await.ok_or(SessionError::McpStartupStopped)?; - let mcp_handle = spawn.handle().clone(); - let mut builder = AgentBuilder::new(provider).max_auto_continues(0); - if let Some(last) = &usage_seed { - builder = builder.resume_usage(last); - } - for prompt in &spec.prompts { - builder = builder.system_prompt(prompt.clone()); + let (servers, stalled_server) = { + let control = self.control.lock().unwrap(); + (control.servers(def), control.stalled_server) + }; + let mut builder = RuntimeBuilder::from_spec(self.cwd.clone(), spec.clone()) + .provider(def.provider.clone()) + .extra_servers(servers) + .elicitation(ElicitationSink::form_and_url(elicitations)) + .resume_usage(usage_seed); + + if !stalled_server { + builder = builder.wait_for_mcp(); } - let (agent_tx, agent_rx, agent_handle) = builder - .tools(mcp_handle, Vec::new()) - .messages(initial_messages) - .spawn() - .await - .map_err(|e| SessionError::Build(CliError::AgentError(e.to_string())))?; - self.control.lock().unwrap().agents.push(agent_tx.clone()); - let (mcp_runtime, event_rx) = spawn.connect_agent(agent_tx.clone()).await.split(); - Ok(AgentRuntime::new(agent_tx, agent_rx, Some(agent_handle), event_rx, mcp_runtime)) + let runtime = builder.build(initial_messages).await?; + self.control.lock().unwrap().agents.push(runtime.agent_tx.clone()); + + Ok(AgentRuntime::new(runtime)) } } struct StubRuntimeFactory { - agent_parts: Mutex>, -} - -struct StubAgentParts { - tx: mpsc::Sender, - rx: mpsc::Receiver, - handle: AgentHandle, - mcp: McpSession, + runtime: Mutex)>>, } #[async_trait::async_trait] @@ -853,33 +907,24 @@ impl RuntimeFactory for StubRuntimeFactory { _spec: &AgentSpec, _initial_messages: Vec, _usage_seed: Option, + elicitations: mpsc::Sender, ) -> Result { - let parts = self - .agent_parts - .lock() - .expect("stub agent parts lock is healthy") - .take() - .expect("stub runtime spawned more than once"); - let (mcp_runtime, event_rx) = parts.mcp.connect_agent(parts.tx.clone()).await.split(); - - Ok(AgentRuntime::new(parts.tx, parts.rx, Some(parts.handle), event_rx, mcp_runtime)) - } -} - -fn with_in_memory_server( - builder: McpBuilder, - name: &str, - server: impl Fn() -> Box> + Send + Sync + 'static, -) -> McpBuilder { - let factory: ServerFactory = Box::new(move |_spec, _services| std::future::ready(server()).boxed()); - let transport = McpTransport::InMemory { - spec: InMemoryServerSpec { factory: name.to_string(), args: Vec::new(), input: None }, - }; - builder.register_in_memory_server(name, factory).with_servers(vec![McpServer::new( - name, - transport, - ToolExposure::ModelVisible, - )]) + let (runtime, mut stub_elicitations) = + self.runtime.lock().expect("stub runtime lock is healthy").take().expect("stub runtime spawned once"); + tokio::spawn(async move { + while let Some(elicitation) = stub_elicitations.recv().await { + if elicitations.send(elicitation).await.is_err() { + return; + } + } + }); + Ok(AgentRuntime::new(runtime)) + } +} + +async fn spawn_ready(builder: McpBuilder) -> (McpCatalog, mpsc::Receiver) { + let (sink, elicitations) = elicitation_sink(); + (ready_catalog(builder.with_elicitation(sink)).await, elicitations) } fn count_prompts(agent_tx: mpsc::Sender, prompts: watch::Sender) -> mpsc::Sender { @@ -960,33 +1005,6 @@ fn fake_agent_spec(name: &str) -> AgentSpec { spec } -fn assert_available_commands(update: SessionUpdate, expected: &[&str], unexpected: &[&str]) { - let SessionUpdate::AvailableCommandsUpdate(commands) = update else { - panic!("expected available commands update"); - }; - let names = commands.available_commands.iter().map(|command| command.name.as_str()).collect::>(); - for name in expected { - assert!(names.contains(name), "expected command /{name} in {names:?}"); - } - for name in unexpected { - assert!(!names.contains(name), "did not expect command /{name} in {names:?}"); - } -} - -fn assert_server_status(notification: McpNotification, expected: &[&str]) { - let McpNotification::ServerStatus { servers } = notification; - let names = servers.iter().map(|server| server.name.as_str()).collect::>(); - for server_name in expected { - assert!(names.contains(server_name), "expected server {server_name} in {names:?}"); - } -} - -fn assert_server_status_exact(notification: McpNotification, expected: &[&str]) { - let McpNotification::ServerStatus { servers } = notification; - let names = servers.iter().map(|server| server.name.as_str()).collect::>(); - assert_eq!(names, expected); -} - fn conversation_texts(context: &Context) -> Vec { context .messages() diff --git a/crates/aether-cli/src/error.rs b/crates/aether-cli/src/error.rs index 594361b3f..f731ab464 100644 --- a/crates/aether-cli/src/error.rs +++ b/crates/aether-cli/src/error.rs @@ -1,4 +1,7 @@ +use crate::resolve::AgentSelectionError; +use crate::runtime::RuntimeError; use aether_auth::OAuthError; +use aether_core::mcp::ParseError; use aether_project::SettingsError; use aether_telemetry::TelemetryInitError; use std::io; @@ -18,14 +21,16 @@ pub enum CliError { Settings(#[from] SettingsError), #[error("Failed to initialize telemetry: {0}")] Telemetry(#[from] TelemetryInitError), - #[error("Model error: {0}")] - ModelError(String), - #[error("MCP error: {0}")] - McpError(String), + #[error(transparent)] + AgentSelection(#[from] AgentSelectionError), + #[error("Invalid MCP config: {0}")] + McpConfig(#[from] ParseError), + #[error(transparent)] + Runtime(#[from] RuntimeError), #[error("IO error: {0}")] IoError(#[from] io::Error), - #[error("Agent error: {0}")] - AgentError(String), + #[error("Agent stopped before receiving the prompt")] + AgentStopped, #[error("Credential store error: {0}")] CredentialStore(#[from] OAuthError), } diff --git a/crates/aether-cli/src/headless/mod.rs b/crates/aether-cli/src/headless/mod.rs index 6e3765281..0e5cbec79 100644 --- a/crates/aether-cli/src/headless/mod.rs +++ b/crates/aether-cli/src/headless/mod.rs @@ -1,12 +1,12 @@ pub mod error; pub mod run; -use aether_core::agent_spec::{AgentSpec, McpConfigSource}; +use aether_core::agent_spec::AgentSpec; +use aether_core::mcp::McpConfig; use aether_project::{AetherSettings, AgentCatalog, TelemetrySettings}; use aether_telemetry::AgentTraceContext; use error::CliError; use llm::{ProviderConnectionOverride, ProviderConnectionOverrides}; -use mcp_utils::client::McpConfig; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; @@ -18,7 +18,7 @@ use crate::mcp_config_args::McpConfigArgs; use crate::output::OutputFormat; use crate::prompt::prompt_or_stdin; use crate::provider_connection_args::ProviderConnectionArgs; -use crate::resolve::{AgentSelectionError, InitialSessionSelection, resolve_agent_from_settings}; +use crate::resolve::{InitialSessionSelection, resolve_agent_from_settings}; use crate::settings_args::SettingsSourceArgs; use aether_auth::OAuthCredentialStorage; use std::sync::Arc; @@ -52,7 +52,7 @@ pub enum CliEventKind { pub struct RunConfig { pub prompt: String, pub cwd: PathBuf, - pub mcp_config_sources: Vec, + pub mcp_config: McpConfig, pub spec: AgentSpec, pub agent_catalog: AgentCatalog, pub system_prompt: Option, @@ -167,14 +167,13 @@ impl RunConfig { return Err(CliError::ConflictingArgs("Cannot specify both --agent and --model".to_string())); } }; - let resolved = resolve_agent_from_settings(&cwd, settings, provider_connections, &selection) - .map_err(map_selection_error)?; - let mcp_config_sources = args.mcp_config.sources(&cwd); + let resolved = resolve_agent_from_settings(&cwd, settings, provider_connections, &selection)?; + let mcp_config = args.mcp_config.load(&cwd)?; Ok(Self { prompt, cwd, - mcp_config_sources, + mcp_config, spec: resolved.spec, agent_catalog: resolved.catalog, system_prompt: args.system_prompt, @@ -203,19 +202,11 @@ impl RunConfig { return Err(CliError::ConflictingArgs("Cannot specify both --agent and --model".to_string())); } }; - let resolved = resolve_agent_from_settings(&cwd, settings, provider_connections, &selection) - .map_err(map_selection_error)?; - let mcp_config_sources = options - .mcp_config - .map(|config| serde_json::to_string(&config).expect("mcp config serialize")) - .map(McpConfigSource::Json) - .into_iter() - .collect(); - + let resolved = resolve_agent_from_settings(&cwd, settings, provider_connections, &selection)?; Ok(Self { prompt, cwd, - mcp_config_sources, + mcp_config: options.mcp_config.unwrap_or_default(), spec: resolved.spec, agent_catalog: resolved.catalog, system_prompt: options.system_prompt, @@ -233,11 +224,3 @@ fn resolve_prompt(args: &HeadlessArgs) -> Result { let explicit = (!args.prompt.is_empty()).then(|| args.prompt.join(" ")); prompt_or_stdin(explicit).map_err(CliError::IoError)?.ok_or(CliError::NoPrompt) } - -fn map_selection_error(error: AgentSelectionError) -> CliError { - match error { - AgentSelectionError::Settings(error) => CliError::Settings(error), - AgentSelectionError::Agent(error) => CliError::AgentError(error.to_string()), - AgentSelectionError::Model(error) => CliError::ModelError(error), - } -} diff --git a/crates/aether-cli/src/headless/run.rs b/crates/aether-cli/src/headless/run.rs index 96991d11f..70899eb18 100644 --- a/crates/aether-cli/src/headless/run.rs +++ b/crates/aether-cli/src/headless/run.rs @@ -2,13 +2,11 @@ use aether_core::core::{AgentDeps, Prompt, recv_agent_event}; use aether_core::events::{ AgentEvent, Command, ContextEvent, MessageEvent, ModelEvent, ToolEvent, TurnEvent, TurnOutcome, }; -use aether_core::mcp::McpHandle; use aether_telemetry::TelemetryRuntime; use std::io; use std::process::ExitCode; use std::sync::Arc; use tokio::sync::mpsc; -use tracing::error; use crate::telemetry::build_telemetry_runtime; @@ -16,7 +14,7 @@ use super::error::CliError; use super::{CliEventKind, RunConfig}; use crate::output::{OutputFormat, print_message}; use crate::runtime::RuntimeBuilder; -use crate::slash_commands::{expand_slash_command, parse_slash_command}; +use crate::slash_commands::expand_slash_command_text; pub async fn run(config: RunConfig) -> Result { setup_tracing(config.verbose); @@ -40,45 +38,28 @@ async fn run_agent(config: RunConfig, telemetry: Option>) let deps = AgentDeps::new(config.oauth_credential_store, telemetry.as_ref().map(|runtime| runtime.observer_factory())) .with_agent_registry(registry); - let (agent, _mcp_snapshot) = RuntimeBuilder::from_spec(config.cwd.clone(), spec) - .mcp_sources(config.mcp_config_sources) + let mut agent = RuntimeBuilder::from_spec(config.cwd.clone(), spec) + .mcp_config(config.mcp_config) .agent_deps(deps) - .build_ready(vec![]) + .wait_for_mcp() + .build(Vec::new()) .await?; - let prompt = expand_prompt(agent.mcp_runtime.handle(), config.prompt).await; + let prompt = expand_slash_command_text(&agent.mcp, config.prompt).await; - agent - .agent_tx - .send(Command::text(&prompt)) - .await - .map_err(|e| CliError::AgentError(format!("Failed to send prompt: {e}")))?; - - let exit_code = stream_output(agent.agent_rx, config.output, &config.events).await; - - drop(agent.agent_tx); - agent.agent_handle.await_completion().await; - - Ok(exit_code) -} - -async fn expand_prompt(mcp: &McpHandle, prompt: String) -> String { - let Some(slash_command) = parse_slash_command(&prompt) else { - return prompt; + let outcome = match agent.agent_tx.send(Command::text(&prompt)).await { + Ok(()) => Ok(stream_output(&mut agent.agent_rx, config.output, &config.events).await), + Err(_) => Err(CliError::AgentStopped), }; - match expand_slash_command(mcp, slash_command.command_name, slash_command.args_text).await { - Ok(expanded) => expanded, - Err(error) => { - error!("Failed to expand slash command: {error}"); - prompt - } - } + agent.shutdown().await; + + outcome } -async fn stream_output(mut rx: mpsc::Receiver, format: OutputFormat, events: &[CliEventKind]) -> ExitCode { +async fn stream_output(rx: &mut mpsc::Receiver, format: OutputFormat, events: &[CliEventKind]) -> ExitCode { loop { - let msg = recv_agent_event(&mut rx).await; + let msg = recv_agent_event(rx).await; if should_emit(&msg, events) && let Err(error) = print_message(format, &msg) { @@ -301,27 +282,27 @@ mod tests { #[tokio::test] async fn stream_output_turn_ended_breaks_loop_under_filter() { - let (tx, rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); tx.send(AgentEvent::turn_ended(TurnOutcome::Completed)).await.unwrap(); let filter = vec![CliEventKind::ToolCall]; - let code = stream_output(rx, OutputFormat::Text, &filter).await; + let code = stream_output(&mut rx, OutputFormat::Text, &filter).await; assert_eq!(code, ExitCode::SUCCESS); } #[tokio::test] async fn stream_output_failed_turn_exits_with_failure() { - let (tx, rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); tx.send(AgentEvent::turn_ended(TurnOutcome::failed("boom"))).await.unwrap(); - let code = stream_output(rx, OutputFormat::Text, &[]).await; + let code = stream_output(&mut rx, OutputFormat::Text, &[]).await; assert_eq!(code, ExitCode::FAILURE); } #[tokio::test] async fn stream_output_channel_closed_before_turn_ended_exits_with_failure() { - let (tx, rx) = mpsc::channel(4); + let (tx, mut rx) = mpsc::channel(4); tx.send(AgentEvent::Turn(TurnEvent::Started { content: vec![] })).await.unwrap(); drop(tx); - let code = stream_output(rx, OutputFormat::Text, &[]).await; + let code = stream_output(&mut rx, OutputFormat::Text, &[]).await; assert_eq!(code, ExitCode::FAILURE); } diff --git a/crates/aether-cli/src/init/build_settings.rs b/crates/aether-cli/src/init/build_settings.rs index 4d28b9a99..f81e34874 100644 --- a/crates/aether-cli/src/init/build_settings.rs +++ b/crates/aether-cli/src/init/build_settings.rs @@ -1,10 +1,10 @@ use super::InitScope; use super::harness::HarnessIntegration; use super::recommendations::{ProviderRecommendations, recommended_for_provider}; +use aether_core::mcp::{InMemoryServerConfig, InMemoryType, McpConfig, McpServerConfig}; +use aether_core::mcp::{ToolExposure, ToolFilter, ToolMatcher}; use aether_project::{AetherSettings, AgentConfig, McpSourceSpec, PromptSource}; use llm::{ReasoningEffort, catalog::Provider}; -use mcp_utils::client::{InMemoryServerConfig, InMemoryType, McpServerConfig, ToolExposure}; -use mcp_utils::client::{ToolFilter, ToolMatcher}; const SYSTEM_PATH: &str = "SYSTEM.md"; const SYSTEM_MD: &str = include_str!("templates/SYSTEM.md"); @@ -256,13 +256,12 @@ fn mcps(servers: Vec<(&str, Vec)>) -> McpSourceSpec { McpServerConfig::InMemory(InMemoryServerConfig { type_: InMemoryType::InMemory, args, - input: None, defer_tools: ToolExposure::ModelVisible, }), ) }) .collect(); - McpSourceSpec::Inline { servers } + McpSourceSpec::Inline(McpConfig::new(servers)) } fn plan_tools() -> ToolFilter { diff --git a/crates/aether-cli/src/lib.rs b/crates/aether-cli/src/lib.rs index 824fb141b..17f497188 100644 --- a/crates/aether-cli/src/lib.rs +++ b/crates/aether-cli/src/lib.rs @@ -13,7 +13,7 @@ pub mod output; pub(crate) mod prompt; pub mod provider_connection_args; pub mod resolve; -pub mod runtime; +pub(crate) mod runtime; pub mod sandbox; pub mod settings; pub mod settings_args; @@ -21,5 +21,3 @@ pub mod show_prompt; pub(crate) mod slash_commands; pub(crate) mod telemetry; pub mod workspace; - -pub use acp::map_mcp_prompt_to_available_command; diff --git a/crates/aether-cli/src/mcp_command.rs b/crates/aether-cli/src/mcp_command.rs index 34ce46589..e4423f104 100644 --- a/crates/aether-cli/src/mcp_command.rs +++ b/crates/aether-cli/src/mcp_command.rs @@ -1,17 +1,13 @@ +use aether_core::mcp::{AETHER_MCP_IPC_SOCKET, DeferredServer, DeferredToolsClient, DeferredToolsError}; use clap::{ArgAction, Args}; -use mcp_utils::ServiceExt; -use mcp_utils::tool_gateway::{AETHER_MCP_IPC_SOCKET, LIST_SERVERS_TOOL, UnixSocketPath, connect}; -use rmcp::model::{CallToolRequestParams, CallToolResponse, CallToolResult, Tool}; -use serde::Deserialize; +use mcp_utils::client::ToolCallError; +use mcp_utils::model::{CallToolResult, Tool}; use serde_json::{Map, Value}; use std::env::var_os; -use std::fmt::Display; -use std::future::Future; use std::io::{self, IsTerminal, Read}; use std::path::PathBuf; use std::time::Duration; -const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); const DEFAULT_CALL_TIMEOUT_SECONDS: u64 = 600; #[derive(Debug, Args)] @@ -45,14 +41,14 @@ pub enum McpCommandError { Usage(String), #[error("`aether mcp` requires the inherited {AETHER_MCP_IPC_SOCKET} from an active Aether session")] SessionUnavailable, - #[error("invalid {AETHER_MCP_IPC_SOCKET}: {0}")] - InvalidSocket(String), - #[error("failed to connect to the active Aether session: {0}")] - Connect(String), + #[error("invalid {AETHER_MCP_IPC_SOCKET}: socket path must be absolute")] + RelativeSocket, + #[error(transparent)] + Session(#[from] DeferredToolsError), #[error("MCP request timed out after {0} seconds")] Timeout(u64), - #[error("MCP request failed: {0}")] - Request(String), + #[error("MCP tool call failed: {0}")] + Call(#[source] ToolCallError), #[error("deferred tool returned an error: {0}")] Tool(String), #[error("failed to read JSON from stdin: {0}")] @@ -74,21 +70,15 @@ enum HelpLevel { Tool { server: String, tool: String }, } -#[derive(Deserialize)] -struct ServerSummary { - name: String, - description: String, -} - impl McpCommandError { pub fn exit_code(&self) -> u8 { match self { Self::Usage(_) | Self::Stdin(_) => 2, Self::SessionUnavailable - | Self::InvalidSocket(_) - | Self::Connect(_) + | Self::RelativeSocket + | Self::Session(_) | Self::Timeout(_) - | Self::Request(_) + | Self::Call(_) | Self::Tool(_) | Self::Output(_) => 1, } @@ -97,9 +87,7 @@ impl McpCommandError { pub async fn run(args: McpArgs) -> Result<(), McpCommandError> { let request = Request::try_from(args)?; - let socket = inherited_socket()?; - let transport = connect(socket.path()).await.map_err(|error| McpCommandError::Connect(error.to_string()))?; - let client = ().serve(transport).await.map_err(|error| McpCommandError::Connect(error.to_string()))?; + let client = DeferredToolsClient::connect(inherited_socket()?).await?; execute_request(&client, request).await } @@ -127,9 +115,9 @@ impl TryFrom for Request { } } -fn inherited_socket() -> Result { - let path = var_os(AETHER_MCP_IPC_SOCKET).ok_or(McpCommandError::SessionUnavailable)?; - UnixSocketPath::from_path(PathBuf::from(path)).map_err(|error| McpCommandError::InvalidSocket(error.to_string())) +fn inherited_socket() -> Result { + let path = PathBuf::from(var_os(AETHER_MCP_IPC_SOCKET).ok_or(McpCommandError::SessionUnavailable)?); + if path.is_absolute() { Ok(path) } else { Err(McpCommandError::RelativeSocket) } } fn read_stdin() -> Result, McpCommandError> { @@ -150,10 +138,7 @@ fn parse_json_object(input: Option<&str>) -> Result, McpComma }) } -async fn execute_request( - client: &rmcp::service::RunningService, - request: Request, -) -> Result<(), McpCommandError> { +async fn execute_request(client: &DeferredToolsClient, request: Request) -> Result<(), McpCommandError> { match request { Request::Help(HelpLevel::Servers) => show_servers_help(client).await, Request::Help(HelpLevel::Server(server)) => show_server_help(client, &server).await, @@ -164,58 +149,38 @@ async fn execute_request( } } -async fn show_servers_help( - client: &rmcp::service::RunningService, -) -> Result<(), McpCommandError> { - let result = timed(DISCOVERY_TIMEOUT, client.call_tool_once(CallToolRequestParams::new(LIST_SERVERS_TOOL))).await?; - let result = complete(result)?; - let servers: Vec = serde_json::from_value( - result - .structured_content - .ok_or_else(|| McpCommandError::Request("server discovery returned no JSON".into()))?, - ) - .map_err(|error| McpCommandError::Request(format!("invalid server discovery response: {error}")))?; - print_servers_help(&servers); +async fn show_servers_help(client: &DeferredToolsClient) -> Result<(), McpCommandError> { + print_servers_help(&client.servers().await?); Ok(()) } -async fn show_server_help( - client: &rmcp::service::RunningService, - server: &str, -) -> Result<(), McpCommandError> { - let tools = list_tools(client).await?; - let tools = tools_for_server(&tools, server); - if tools.is_empty() { - return Err(McpCommandError::Usage(format!("unknown deferred server `{server}`"))); - } - print_server_help(server, tools); +async fn show_server_help(client: &DeferredToolsClient, server: &str) -> Result<(), McpCommandError> { + let servers = client.servers().await?; + let deferred = servers + .iter() + .find(|deferred| deferred.name == server) + .ok_or_else(|| McpCommandError::Usage(format!("unknown deferred server `{server}`")))?; + print_server_help(deferred); Ok(()) } -async fn show_tool_help( - client: &rmcp::service::RunningService, - server: &str, - tool: &str, -) -> Result<(), McpCommandError> { - let tools = list_tools(client).await?; - let namespaced = format!("{server}__{tool}"); - let definition = tools - .into_iter() - .find(|definition| definition.name == namespaced) +async fn show_tool_help(client: &DeferredToolsClient, server: &str, tool: &str) -> Result<(), McpCommandError> { + let definition = client + .tool(server, tool) + .await? .ok_or_else(|| McpCommandError::Usage(format!("unknown deferred tool `{server} {tool}`")))?; print_tool_help(server, tool, &definition) } async fn call_tool( - client: &rmcp::service::RunningService, + client: &DeferredToolsClient, server: &str, tool: &str, json: Map, timeout_seconds: u64, ) -> Result<(), McpCommandError> { - let params = CallToolRequestParams::new(format!("{server}__{tool}")).with_arguments(json); - let result = timed(Duration::from_secs(timeout_seconds), client.call_tool_once(params)).await?; - let result = complete(result)?; + let call = client.call(server, tool, json, Duration::from_secs(timeout_seconds)); + let result = call.result().await.map_err(call_error)?; if result.is_error.unwrap_or(false) { return Err(McpCommandError::Tool(result_json(&result).to_string())); } @@ -223,35 +188,14 @@ async fn call_tool( Ok(()) } -async fn list_tools( - client: &rmcp::service::RunningService, -) -> Result, McpCommandError> { - timed(DISCOVERY_TIMEOUT, client.list_all_tools()).await -} - -async fn timed( - duration: Duration, - future: impl Future>, -) -> Result { - tokio::time::timeout(duration, future) - .await - .map_err(|_| McpCommandError::Timeout(duration.as_secs()))? - .map_err(|error| McpCommandError::Request(error.to_string())) -} - -fn complete(response: CallToolResponse) -> Result { - match response { - CallToolResponse::Complete(result) => Ok(result), - other => Err(McpCommandError::Request(format!("gateway returned an incomplete response: {other:?}"))), +fn call_error(error: ToolCallError) -> McpCommandError { + match error { + ToolCallError::TimedOut(timeout) => McpCommandError::Timeout(timeout.as_secs()), + error => McpCommandError::Call(error), } } -fn tools_for_server<'a>(tools: &'a [Tool], server: &str) -> Vec<&'a Tool> { - let prefix = format!("{server}__"); - tools.iter().filter(|tool| tool.name.starts_with(&prefix)).collect() -} - -fn print_servers_help(servers: &[ServerSummary]) { +fn print_servers_help(servers: &[DeferredServer]) { println!("Discover and call deferred MCP tools through the active Aether session.\n"); println!("Usage: aether mcp --help\n"); println!("Deferred MCP servers:"); @@ -260,13 +204,13 @@ fn print_servers_help(servers: &[ServerSummary]) { } } -fn print_server_help(server: &str, tools: Vec<&Tool>) { - println!("Deferred tools from `{server}`.\n"); - println!("Usage: aether mcp {server} --help\n"); +fn print_server_help(server: &DeferredServer) { + let name = &server.name; + println!("Deferred tools from `{name}`.\n"); + println!("Usage: aether mcp {name} --help\n"); println!("Tools:"); - for tool in tools { - let local_name = tool.name.strip_prefix(&format!("{server}__")).unwrap_or(tool.name.as_ref()); - println!(" {:<20} {}", local_name, tool.description.as_deref().unwrap_or_default()); + for tool in &server.tools { + println!(" {:<20} {}", tool.name, tool.description.as_deref().unwrap_or_default()); } } diff --git a/crates/aether-cli/src/mcp_config_args.rs b/crates/aether-cli/src/mcp_config_args.rs index 0cf060e84..ef65cc477 100644 --- a/crates/aether-cli/src/mcp_config_args.rs +++ b/crates/aether-cli/src/mcp_config_args.rs @@ -1,4 +1,4 @@ -use aether_core::agent_spec::McpConfigSource; +use aether_core::mcp::{McpConfig, ParseError}; use std::path::{Path, PathBuf}; #[derive(Clone, Debug, Default, clap::Args)] @@ -11,12 +11,9 @@ pub struct McpConfigArgs { } impl McpConfigArgs { - pub fn sources(&self, project_root: &Path) -> Vec { - self.mcp_configs - .iter() - .map(|path| project_root.join(path)) - .map(McpConfigSource::model_visible) - .chain(self.mcp_config_jsons.iter().cloned().map(McpConfigSource::Json)) - .collect() + pub fn load(&self, project_root: &Path) -> Result { + let files = self.mcp_configs.iter().map(|path| McpConfig::from_json_file(project_root.join(path))); + let jsons = self.mcp_config_jsons.iter().map(|json| McpConfig::from_json(json)); + files.chain(jsons).collect() } } diff --git a/crates/aether-cli/src/resolve.rs b/crates/aether-cli/src/resolve.rs index 9d13ee47c..b394fe54b 100644 --- a/crates/aether-cli/src/resolve.rs +++ b/crates/aether-cli/src/resolve.rs @@ -33,7 +33,7 @@ pub(crate) struct ResolvedAgentSelection { } #[derive(Debug, Error)] -pub(crate) enum AgentSelectionError { +pub enum AgentSelectionError { #[error(transparent)] Settings(#[from] SettingsError), #[error("{0}")] @@ -76,10 +76,5 @@ pub fn resolve_agent_spec( ) -> Result { let selection = agent_name.map_or(InitialSessionSelection::Default, |name| InitialSessionSelection::Agent(name.to_string())); - resolve_agent_from_catalog(catalog.clone(), &selection).map(|resolved| resolved.spec).map_err(|error| match error { - AgentSelectionError::Settings(error) | AgentSelectionError::Agent(error) => { - crate::error::CliError::AgentError(error.to_string()) - } - AgentSelectionError::Model(error) => crate::error::CliError::ModelError(error), - }) + Ok(resolve_agent_from_catalog(catalog.clone(), &selection)?.spec) } diff --git a/crates/aether-cli/src/runtime.rs b/crates/aether-cli/src/runtime.rs index 38564a0a7..cb36f0811 100644 --- a/crates/aether-cli/src/runtime.rs +++ b/crates/aether-cli/src/runtime.rs @@ -1,23 +1,23 @@ -use crate::error::CliError; -use aether_core::agent_spec::{AgentSpec, McpConfigSource}; -use aether_core::core::{AgentBuilder, AgentDeps, AgentHandle, Prompt}; +use aether_core::agent_spec::AgentSpec; +use aether_core::core::{AgentBuilder, AgentDeps, AgentError, AgentHandle, Prompt, PromptCache}; use aether_core::events::{AgentEvent, Command}; -use aether_core::mcp::McpBuilder; -use aether_core::mcp::mcp; -use aether_core::mcp::{McpHandle, McpRuntime, McpSession}; -use llm::{ChatMessage, SessionUsageEvent, ToolDefinition}; -use mcp_servers::McpBuilderExt; -use mcp_utils::client::{McpClientEvent, McpConnectionDetails, McpServer, OAuthHandlerFactory}; +use aether_core::mcp::{McpCatalog, McpConfig, McpSpawnError, ParseError}; +use llm::{ChatMessage, SessionUsageEvent, StreamingModelProvider, ToolDefinition}; +use mcp_servers::agent_mcp; +use mcp_utils::client::ElicitationSink; +use mcp_utils::gateway::ServerSpec; use std::path::PathBuf; +use std::sync::Arc; +use thiserror::Error; use tokio::sync::mpsc::{Receiver, Sender}; -use tracing::debug; pub struct RuntimeBuilder { cwd: PathBuf, spec: AgentSpec, - mcp_config_sources: Vec, - extra_mcp_servers: Vec, - oauth_applicator: Option McpBuilder + Send>>, + provider: Option>, + extra_mcp_servers: Vec, + elicitation: Option, + wait_for_mcp: bool, agent_deps: AgentDeps, usage_seed: Option, } @@ -26,13 +26,22 @@ pub struct Runtime { pub agent_tx: Sender, pub agent_rx: Receiver, pub agent_handle: AgentHandle, - pub event_rx: Receiver, - pub mcp_runtime: McpRuntime, + pub mcp: McpCatalog, } -pub struct PromptInfo { - pub spec: AgentSpec, - pub tool_definitions: Vec, +pub struct PromptPreview { + pub system_prompt: String, + pub tools: Vec, +} + +#[derive(Debug, Error)] +pub enum RuntimeError { + #[error("Invalid MCP config: {0}")] + McpConfig(#[from] ParseError), + #[error("MCP error: {0}")] + McpSpawn(#[from] McpSpawnError), + #[error("Agent error: {0}")] + Agent(#[from] AgentError), } impl RuntimeBuilder { @@ -40,9 +49,10 @@ impl RuntimeBuilder { Self { cwd, spec, - mcp_config_sources: Vec::new(), + provider: None, extra_mcp_servers: Vec::new(), - oauth_applicator: None, + elicitation: None, + wait_for_mcp: false, agent_deps: AgentDeps::default(), usage_seed: None, } @@ -53,136 +63,87 @@ impl RuntimeBuilder { self } - /// Continue session usage totals from the last persisted usage event. - pub fn resume_usage(mut self, last: SessionUsageEvent) -> Self { - self.usage_seed = Some(last); + #[cfg(any(test, feature = "testing"))] + pub fn provider(mut self, provider: Arc) -> Self { + self.provider = Some(provider); self } - /// Set MCP config source overrides. When non-empty, these completely - /// replace any sources resolved from the agent's `AgentSpec`. - pub fn mcp_sources(mut self, sources: Vec) -> Self { - self.mcp_config_sources = sources; + pub fn resume_usage(mut self, last: Option) -> Self { + self.usage_seed = last; self } - pub fn extra_servers(mut self, servers: Vec) -> Self { - self.extra_mcp_servers = servers; + /// Replaces the agent's own MCP servers when `config` names any. + pub fn mcp_config(mut self, config: McpConfig) -> Self { + if !config.is_empty() { + self.spec.mcp = config; + } self } - pub fn oauth_handler_factory(mut self, factory: OAuthHandlerFactory) -> Self { - self.oauth_applicator = Some(Box::new(|builder| builder.with_oauth_handler_factory(factory))); + pub fn extra_servers(mut self, servers: Vec) -> Self { + self.extra_mcp_servers = servers; self } - pub async fn build( - self, - custom_prompt: Option, - messages: Option>, - ) -> Result { - let deps = self.agent_deps.clone(); - let usage_seed = self.usage_seed.clone(); - let (spec, session) = self.spawn_mcp().await?; - let mcp = session.handle().clone(); - - let (agent_tx, agent_rx, agent_handle) = spawn_agent(&spec, &deps, mcp, Vec::new(), |mut agent_builder| { - if let Some(last) = &usage_seed { - agent_builder = agent_builder.resume_usage(last); - } - if let Some(prompt) = custom_prompt { - agent_builder = agent_builder.system_prompt(prompt); - } - if let Some(msgs) = messages { - agent_builder = agent_builder.messages(msgs); - } - agent_builder - }) - .await?; - let (mcp_runtime, event_rx) = session.connect_agent(agent_tx.clone()).await.split(); - - Ok(Runtime { agent_tx, agent_rx, agent_handle, event_rx, mcp_runtime }) - } - - /// Spawn MCP, block until every server finishes its initial connection, - /// then connect the agent to the session's filtered tools and MCP - /// instructions before returning. Returns the live [`Runtime`] plus the - /// bootstrap snapshot for callers that need the agent ready to use tools on - /// its first turn. - pub async fn build_ready(self, messages: Vec) -> Result<(Runtime, McpConnectionDetails), CliError> { - let deps = self.agent_deps.clone(); - let usage_seed = self.usage_seed.clone(); - let (spec, mut session) = self.spawn_mcp().await?; - let snapshot = session - .block_until_ready() - .await - .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?; - let mcp = session.handle().clone(); - - let (agent_tx, agent_rx, agent_handle) = spawn_agent(&spec, &deps, mcp, Vec::new(), |mut agent_builder| { - if let Some(last) = &usage_seed { - agent_builder = agent_builder.resume_usage(last); - } - agent_builder.messages(messages) - }) - .await?; - let (mcp_runtime, event_rx) = session.connect_agent(agent_tx.clone()).await.split(); - - Ok((Runtime { agent_tx, agent_rx, agent_handle, event_rx, mcp_runtime }, snapshot)) + pub fn elicitation(mut self, sink: ElicitationSink) -> Self { + self.elicitation = Some(sink); + self } - pub async fn build_prompt_info(self) -> Result { - let (spec, mut session) = self.spawn_mcp().await?; - let details = session - .block_until_ready() - .await - .ok_or_else(|| CliError::McpError("MCP bootstrap aborted before completion".to_string()))?; - let filtered_tools = details.tool_definitions(); - Ok(PromptInfo { spec, tool_definitions: filtered_tools }) + pub fn wait_for_mcp(mut self) -> Self { + self.wait_for_mcp = true; + self } - async fn spawn_mcp(self) -> Result<(AgentSpec, McpSession), CliError> { - let deps = self.agent_deps.clone(); - let mut builder = mcp(&self.cwd).with_tool_filter(self.spec.tools.clone()); - - if let Some(apply_oauth) = self.oauth_applicator { - builder = apply_oauth(builder); + pub async fn build(self, messages: Vec) -> Result { + let mcp = self.spawn_mcp()?; + if self.wait_for_mcp { + mcp.ready().await; } - builder = builder.with_agent_deps(deps).with_builtin_servers(); - - if !self.extra_mcp_servers.is_empty() { - builder = builder.with_servers(self.extra_mcp_servers); + let builder = match self.provider.clone() { + Some(provider) => AgentBuilder::from_spec_with_llm(&self.spec, &self.agent_deps, provider), + None => AgentBuilder::from_spec(&self.spec, &self.agent_deps).await?, + }; + let mut builder = builder.mcp(mcp.clone()).messages(messages); + if let Some(last) = &self.usage_seed { + builder = builder.resume_usage(last); } + let (agent_tx, agent_rx, agent_handle) = builder.spawn().await?; - let mcp_config_sources: Vec = if self.mcp_config_sources.is_empty() { - self.spec.mcp_config_sources.clone() - } else { - self.mcp_config_sources - }; + Ok(Runtime { agent_tx, agent_rx, agent_handle, mcp }) + } - if !mcp_config_sources.is_empty() { - debug!("Loading MCP configs from: {:?}", mcp_config_sources); - builder = - builder.from_mcp_config_sources(&mcp_config_sources).map_err(|e| CliError::McpError(e.to_string()))?; - } + pub async fn preview(self) -> Result { + let mcp = self.spawn_mcp()?; + mcp.ready().await; + let mut prompts = self.spec.prompts.clone(); + prompts.push(Prompt::McpInstructions(mcp.instructions())); + let system_prompt = PromptCache::new(prompts).render().await; + let tools = mcp.tool_definitions(); + mcp.shutdown().await; + Ok(PromptPreview { system_prompt: system_prompt?, tools }) + } - let spawn = builder.spawn().await.map_err(|e| CliError::McpError(e.to_string()))?; - Ok((self.spec, spawn)) + fn spawn_mcp(&self) -> Result { + let mut builder = agent_mcp(&self.cwd, self.agent_deps.clone(), self.spec.tools.clone()) + .with_servers(self.extra_mcp_servers.clone()) + .with_config(self.spec.mcp.clone())?; + if let Some(sink) = self.elicitation.clone() { + builder = builder.with_elicitation(sink); + } + Ok(builder.spawn()?) } } -async fn spawn_agent( - spec: &AgentSpec, - deps: &AgentDeps, - mcp: McpHandle, - tool_definitions: Vec, - configure: impl FnOnce(AgentBuilder) -> AgentBuilder, -) -> Result<(Sender, Receiver, AgentHandle), CliError> { - let builder = AgentBuilder::from_spec(spec, vec![], deps) - .await - .map_err(|error| CliError::AgentError(error.to_string()))? - .tools(mcp, tool_definitions); - - configure(builder).spawn().await.map_err(|error| CliError::AgentError(error.to_string())) +impl Runtime { + pub async fn shutdown(self) { + let Self { agent_tx, agent_rx, agent_handle, mcp } = self; + drop((agent_tx, agent_rx)); + agent_handle.abort(); + agent_handle.await_completion().await; + mcp.shutdown().await; + } } diff --git a/crates/aether-cli/src/show_prompt/run.rs b/crates/aether-cli/src/show_prompt/run.rs index 232150829..df3cefded 100644 --- a/crates/aether-cli/src/show_prompt/run.rs +++ b/crates/aether-cli/src/show_prompt/run.rs @@ -3,49 +3,43 @@ use std::collections::BTreeMap; use super::PromptArgs; use crate::error::CliError; use crate::resolve::resolve_agent_spec; -use crate::runtime::RuntimeBuilder; +use crate::runtime::{PromptPreview, RuntimeBuilder}; use aether_core::core::{AgentDeps, Prompt}; use llm::ToolDefinition; use serde_json::Value; pub async fn run_prompt(args: PromptArgs) -> Result<(), CliError> { let cwd = args.cwd.canonicalize().map_err(CliError::IoError)?; - let catalog = args.settings_source.load_agent_catalog(&cwd).map_err(|e| CliError::AgentError(e.to_string()))?; - let spec = resolve_agent_spec(&catalog, args.agent.as_deref())?; + let catalog = args.settings_source.load_agent_catalog(&cwd)?; + let mut spec = resolve_agent_spec(&catalog, args.agent.as_deref())?; + if let Some(custom) = args.system_prompt.as_deref() { + spec.prompts.push(Prompt::text(custom)); + } let registry = catalog.registry().clone(); - let info = RuntimeBuilder::from_spec(cwd.clone(), spec) + let PromptPreview { system_prompt, tools: tool_definitions } = RuntimeBuilder::from_spec(cwd.clone(), spec) .agent_deps(AgentDeps::default().with_agent_registry(registry)) - .mcp_sources(args.mcp_config.sources(&cwd)) - .build_prompt_info() + .mcp_config(args.mcp_config.load(&cwd)?) + .preview() .await?; - let system_prompt = build_prompt(&info.spec.prompts, args.system_prompt.as_deref()).await?; - let tools_output = build_tools(&info.tool_definitions); + let tools_output = build_tools(&tool_definitions); println!("{system_prompt}"); if !tools_output.is_empty() { println!(); - println!("--- Tools ({} tools) ---", info.tool_definitions.len()); + println!("--- Tools ({} tools) ---", tool_definitions.len()); println!(); println!("{tools_output}"); } println!(); - println!("{}", format_stats(system_prompt.len(), tools_output.len(), info.tool_definitions.len())); + println!("{}", format_stats(system_prompt.len(), tools_output.len(), tool_definitions.len())); Ok(()) } -pub async fn build_prompt(prompts: &[Prompt], custom: Option<&str>) -> Result { - let mut prompts = prompts.to_vec(); - if let Some(custom) = custom { - prompts.push(Prompt::text(custom)); - } - Prompt::build_all(&prompts).await.map_err(|e| CliError::AgentError(e.to_string())) -} - pub fn build_tools(tools: &[ToolDefinition]) -> String { if tools.is_empty() { return String::new(); diff --git a/crates/aether-cli/src/slash_commands.rs b/crates/aether-cli/src/slash_commands.rs index 62852ba4c..fb43f23bd 100644 --- a/crates/aether-cli/src/slash_commands.rs +++ b/crates/aether-cli/src/slash_commands.rs @@ -1,25 +1,66 @@ -use aether_core::mcp::McpHandle; -use agent_client_protocol::schema::v2::AvailableCommand; -use rmcp::model::{ContentBlock, GetPromptResult, Prompt as McpPrompt}; +use aether_core::mcp::McpCatalog; +use llm::ContentBlock as MessageContent; +use mcp_utils::gateway::{GatewayError, ServerPrompt}; +use mcp_utils::model::{ContentBlock, GetPromptResult}; use std::collections::HashSet; use thiserror::Error; +use tracing::{debug, error, info}; -pub(crate) struct SlashCommand<'a> { - pub(crate) command_name: &'a str, - pub(crate) args_text: &'a str, +pub(crate) async fn expand_slash_command_in_content( + mcp: &McpCatalog, + mut content: Vec, +) -> Vec { + if let Some(MessageContent::Text { text }) = content.first() { + let expanded = expand_slash_command_text(mcp, text.clone()).await; + content[0] = MessageContent::text(expanded); + } + content +} + +pub(crate) async fn expand_slash_command_text(mcp: &McpCatalog, text: String) -> String { + let Some(slash_command) = parse_slash_command(&text) else { + return text; + }; + + match expand_slash_command(mcp, slash_command.command_name, slash_command.args_text).await { + Ok(expanded) => { + info!("Expanded slash command -> {} chars", expanded.len()); + expanded + } + Err(SlashCommandError::NotFound(command)) => { + debug!(command, "Passing an unknown slash command through unexpanded"); + text + } + Err(error) => { + error!("Failed to expand slash command: {error}"); + text + } + } +} + +pub(crate) async fn prompt_commands(mcp: &McpCatalog) -> Vec { + let mut seen_names = HashSet::new(); + let mut prompts = mcp.list_prompts().await; + prompts.retain(|listed| seen_names.insert(listed.prompt.name.clone())); + prompts +} + +struct SlashCommand<'a> { + command_name: &'a str, + args_text: &'a str, } #[derive(Error, Debug)] -pub(crate) enum SlashCommandError { +enum SlashCommandError { #[error("MCP operation failed: {0}")] - McpOperation(String), + McpOperation(#[from] GatewayError), #[error("slash command '/{0}' not found")] NotFound(String), #[error("prompt result contains no text content")] NoTextContent, } -pub(crate) fn parse_slash_command(text: &str) -> Option> { +fn parse_slash_command(text: &str) -> Option> { let slash_command_text = text.strip_prefix('/')?; let (command_name, args_text) = if let Some(space_idx) = slash_command_text.find(char::is_whitespace) { @@ -32,43 +73,21 @@ pub(crate) fn parse_slash_command(text: &str) -> Option> { Some(SlashCommand { command_name, args_text }) } -pub(crate) fn dedupe_commands_by_name(commands: Vec) -> Vec { - let mut seen_names = HashSet::new(); - commands.into_iter().filter(|command| seen_names.insert(command.name.clone())).collect() -} - -pub(crate) async fn expand_slash_command( - mcp: &McpHandle, +async fn expand_slash_command( + mcp: &McpCatalog, command_name: &str, args_text: &str, ) -> Result { + let command = prompt_commands(mcp) + .await + .into_iter() + .find(|listed| listed.prompt.name == command_name) + .ok_or_else(|| SlashCommandError::NotFound(command_name.to_string()))?; let arguments = parse_slash_command_arguments(args_text); - let prompts = list_prompts(mcp).await?; - let prompt_name = find_prompt_name(&prompts, command_name)?; - let prompt_result = get_prompt(mcp, prompt_name, arguments).await?; + let prompt_result = mcp.get_prompt(&command.server, &command.prompt.name, arguments).await?; prompt_result_text(&prompt_result) } -pub(crate) async fn list_prompts(mcp: &McpHandle) -> Result, SlashCommandError> { - mcp.list_prompts().await.map_err(|error| SlashCommandError::McpOperation(error.to_string())) -} - -async fn get_prompt( - mcp: &McpHandle, - name: String, - arguments: Option>, -) -> Result { - mcp.get_prompt(&name, arguments).await.map_err(|error| SlashCommandError::McpOperation(error.to_string())) -} - -fn find_prompt_name(prompts: &[McpPrompt], command_name: &str) -> Result { - prompts - .iter() - .find(|p| p.name.split("__").last().unwrap_or("") == command_name) - .map(|prompt| prompt.name.clone()) - .ok_or_else(|| SlashCommandError::NotFound(command_name.to_string())) -} - fn prompt_result_text(prompt_result: &GetPromptResult) -> Result { prompt_result .messages @@ -80,40 +99,14 @@ fn prompt_result_text(prompt_result: &GetPromptResult) -> Result Option> { +fn parse_slash_command_arguments(args_text: &str) -> serde_json::Map { + let mut arg_map = serde_json::Map::new(); if args_text.is_empty() { - None - } else { - let mut arg_map = serde_json::Map::new(); - arg_map.insert("ARGUMENTS".to_string(), serde_json::Value::String(args_text.to_string())); - for (i, arg) in args_text.split_whitespace().enumerate() { - arg_map.insert((i + 1).to_string(), serde_json::Value::String(arg.to_string())); - } - Some(arg_map) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_argument_parsing() { - let arg_map = parse_slash_command_arguments("do a thing that has spaces").expect("Expected Some"); - let expected = serde_json::Map::from_iter([ - ("ARGUMENTS".to_string(), serde_json::Value::String("do a thing that has spaces".to_string())), - ("1".to_string(), serde_json::Value::String("do".to_string())), - ("2".to_string(), serde_json::Value::String("a".to_string())), - ("3".to_string(), serde_json::Value::String("thing".to_string())), - ("4".to_string(), serde_json::Value::String("that".to_string())), - ("5".to_string(), serde_json::Value::String("has".to_string())), - ("6".to_string(), serde_json::Value::String("spaces".to_string())), - ]); - assert_eq!(arg_map, expected); + return arg_map; } - - #[test] - fn test_empty_arguments_returns_none() { - assert!(parse_slash_command_arguments("").is_none()); + arg_map.insert("ARGUMENTS".to_string(), serde_json::Value::String(args_text.to_string())); + for (i, arg) in args_text.split_whitespace().enumerate() { + arg_map.insert((i + 1).to_string(), serde_json::Value::String(arg.to_string())); } + arg_map } diff --git a/crates/aether-cli/tests/integration/acp_agent_switching.rs b/crates/aether-cli/tests/integration/acp_agent_switching.rs index b5f67908d..ad986dee7 100644 --- a/crates/aether-cli/tests/integration/acp_agent_switching.rs +++ b/crates/aether-cli/tests/integration/acp_agent_switching.rs @@ -154,15 +154,15 @@ async fn mode_change_applies_at_next_prompt_boundary() { #[tokio::test(flavor = "current_thread")] async fn prompt_sent_before_a_deferred_mode_switch_is_validated_against_the_agent_that_runs_it() -> TestResult { AcpTestHarness::run(|mut harness| async move { - let (expansion_started, release_expansion) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let fake = harness.insert_agent_switching_session().await; harness.expect_mcp_server_status(&["planner-mcp"]).await; let in_flight = send_prompt(&harness, &fake, "/plan"); - expansion_started.notified().await; + expansion.reached().await; select_coder(&harness, &fake).await; let image = send_content(&harness, &fake, vec![ContentBlock::Image(ImageContent::new("aW1n", "image/png"))]); - release_expansion.notify_one(); + expansion.open(); in_flight.await?; image.await?; diff --git a/crates/aether-cli/tests/integration/acp_mcp_sign_in.rs b/crates/aether-cli/tests/integration/acp_mcp_sign_in.rs new file mode 100644 index 000000000..455958c44 --- /dev/null +++ b/crates/aether-cli/tests/integration/acp_mcp_sign_in.rs @@ -0,0 +1,46 @@ +use acp_utils::notifications::McpRequest; +use aether_cli::acp::testing::AcpTestHarness; +use agent_client_protocol::schema::v2::ElicitationMode; +use mcp_utils::testing::{FakeAuthServer, FakeMcpServer, complete_authorization}; +use utils::mcp_status::{McpServerAuthCapability, McpServerStatus}; + +#[tokio::test(flavor = "current_thread")] +async fn signing_in_to_an_mcp_server_runs_through_the_client_and_reports_completion() { + AcpTestHarness::run(|mut harness| async move { + let auth = FakeAuthServer::start(FakeMcpServer::math()).await; + harness.add_mcp_server(auth.spec("linear")); + harness + .peer + .queue_elicitation_response(serde_json::from_value(serde_json::json!({"action": "accept"})).unwrap()); + let session = harness.insert_agent_switching_session().await; + let needs_oauth = wait_for_linear(&mut harness, |status| *status == McpServerStatus::NeedsOAuth).await; + assert_eq!(needs_oauth, McpServerAuthCapability::OAuth); + + let session_id = session.session_id().0.to_string(); + harness + .client_cx + .send_notification(McpRequest::Authenticate { session_id, server_name: "linear".to_string() }) + .unwrap(); + + let request = harness.peer.next_elicitation_request().await; + let ElicitationMode::Url(prompt) = request.mode else { panic!("sign-in asks the client to open a URL") }; + complete_authorization(&prompt.url).await; + let completion = harness.peer.next_elicitation_completion().await; + assert_eq!(completion.elicitation_id, prompt.elicitation_id); + wait_for_linear(&mut harness, |status| *status == McpServerStatus::Connected { tool_count: 2 }).await; + harness.shutdown().await; + }) + .await; +} + +async fn wait_for_linear( + harness: &mut AcpTestHarness, + reached: impl Fn(&McpServerStatus) -> bool, +) -> McpServerAuthCapability { + loop { + let statuses = harness.next_mcp_server_statuses().await; + if let Some(linear) = statuses.iter().find(|entry| entry.name == "linear" && reached(&entry.status)) { + return linear.auth_capability; + } + } +} diff --git a/crates/aether-cli/tests/integration/acp_prompt_lifecycle.rs b/crates/aether-cli/tests/integration/acp_prompt_lifecycle.rs index bfd080895..81ebb5776 100644 --- a/crates/aether-cli/tests/integration/acp_prompt_lifecycle.rs +++ b/crates/aether-cli/tests/integration/acp_prompt_lifecycle.rs @@ -11,15 +11,16 @@ use llm::LlmResponse; use llm::testing::{FakeLlmProvider, llm_response}; use std::sync::Arc; use tokio::{sync::Notify, task::LocalSet}; +use utils::mcp_status::{McpServerStatus, McpServerStatusEntry}; #[tokio::test(flavor = "current_thread")] async fn cancel_during_mcp_prompt_expansion_does_not_wait_for_the_server() { AcpTestHarness::run(|mut harness| async move { - let (started, _release) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let session = harness.insert_agent_switching_session().await; let id = session.session_id().clone(); let prompt = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])).block_task(); - started.notified().await; + expansion.reached().await; harness.client_cx.send_notification(CancelSessionNotification::new(id.clone())).unwrap(); let error = prompt.await.expect_err("a prompt cancelled before insertion errors"); assert_eq!(error.code, ErrorCode::RequestCancelled); @@ -33,11 +34,11 @@ async fn cancel_during_mcp_prompt_expansion_does_not_wait_for_the_server() { #[tokio::test(flavor = "current_thread")] async fn cancel_during_prompt_expansion_reports_no_turn() -> TestResult { AcpTestHarness::run(|mut harness| async move { - let (started, _release) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let session = harness.insert_agent_switching_session().await; let id = session.session_id().clone(); let prompt = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])).block_task(); - started.notified().await; + expansion.reached().await; harness.client_cx.send_notification(CancelSessionNotification::new(id.clone()))?; prompt.await.expect_err("a prompt cancelled before insertion errors"); @@ -208,6 +209,23 @@ async fn acceptance_precedes_streaming_and_idle_completes_the_turn() { .await; } +#[tokio::test(flavor = "current_thread")] +async fn mcp_status_is_reported_only_when_a_tool_refresh_changes_it() -> TestResult { + AcpTestHarness::run(|mut harness| async move { + let task = FakeBackgroundTask::new(); + let id = SessionId::new("mcp-status-refresh"); + harness.insert_background_task_session(background_task_provider(), &task, id).await; + let initial = tasks_tool_count(&harness.next_mcp_server_statuses().await).expect("tasks server is connected"); + + task.announce_unchanged_tools().await; + task.add_tool("extra").await; + + assert_eq!(tasks_tool_count(&harness.next_mcp_server_statuses().await), Some(initial + 1)); + Ok(()) + }) + .await +} + #[tokio::test(flavor = "current_thread")] async fn background_task_outcome_turn_is_framed_by_running_and_idle() -> TestResult { AcpTestHarness::run(|mut harness| async move { @@ -255,7 +273,7 @@ async fn prompt_during_background_task_outcome_turn_is_inserted_after_its_reply( #[tokio::test(flavor = "current_thread")] async fn prompt_prepared_during_background_task_outcome_turn_is_inserted_after_its_reply() -> TestResult { AcpTestHarness::run(|mut harness| async move { - let (expansion_started, release_expansion) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let task = FakeBackgroundTask::new(); let release = Arc::new(Notify::new()); let provider = background_task_provider().pause_turn_after(2, 1, release.clone()); @@ -263,13 +281,13 @@ async fn prompt_prepared_during_background_task_outcome_turn_is_inserted_after_i harness.insert_background_task_session(provider, &task, id.clone()).await; start_background_task(&mut harness, &id).await?; let prompt = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])).block_task(); - expansion_started.notified().await; + expansion.reached().await; task.complete("task output"); let updates = harness.updates_until(&id, |update| matches!(update, SessionUpdate::AgentMessageChunk(_))).await; assert!(updates.iter().any(is_running), "task outcome turn must report running before replying: {updates:?}"); let mut delivered = harness.watch_agent_prompts(); - release_expansion.notify_one(); + expansion.open(); delivered.changed().await?; release.notify_one(); let prompt = prompt.await?; @@ -317,14 +335,14 @@ async fn queued_prompt_is_resumed_where_it_was_inserted() -> TestResult { #[tokio::test(flavor = "current_thread")] async fn prompts_sent_together_are_inserted_in_submission_order() -> TestResult { AcpTestHarness::run(|mut harness| async move { - let (expansion_started, release_expansion) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let session = harness.insert_agent_switching_session().await; let id = session.session_id().clone(); let slow = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])).block_task(); - expansion_started.notified().await; + expansion.reached().await; let fast = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["right after".into()])).block_task(); - release_expansion.notify_one(); + expansion.open(); let slow = slow.await?; let fast = fast.await?; @@ -381,6 +399,13 @@ async fn prompt_sent_after_cancel_starts_a_new_turn() -> TestResult { .await } +fn tasks_tool_count(statuses: &[McpServerStatusEntry]) -> Option { + statuses.iter().find(|entry| entry.name == "tasks").and_then(|entry| match entry.status { + McpServerStatus::Connected { tool_count } => Some(tool_count), + _ => None, + }) +} + fn background_task_provider() -> FakeLlmProvider { FakeLlmProvider::new(vec![ llm_response().tool_call("start-call", FakeBackgroundTask::TOOL, &["{}"]).build(), diff --git a/crates/aether-cli/tests/integration/acp_remote.rs b/crates/aether-cli/tests/integration/acp_remote.rs index 1e31e362b..d9697b445 100644 --- a/crates/aether-cli/tests/integration/acp_remote.rs +++ b/crates/aether-cli/tests/integration/acp_remote.rs @@ -6,8 +6,9 @@ use aether_core::core::agent; use aether_core::events::{AgentEvent, MessageEvent, TurnEvent, TurnOutcome}; use aether_sessions::{SessionEvent, UserEvent}; use agent_client_protocol::schema::v2::{ - AbsolutePath, CloseSessionRequest, ContentBlock, NewSessionRequest, PromptRequest, ReplayFrom, ReplayFromStart, - ResumeSessionRequest, SessionId, SessionUpdate, SetSessionConfigOptionRequest, StateUpdate, StopReason, + AbsolutePath, CloseSessionRequest, ContentBlock, ElicitationMode, ElicitationScope, NewSessionRequest, + PromptRequest, ReplayFrom, ReplayFromStart, ResumeSessionRequest, SessionId, SessionUpdate, + SetSessionConfigOptionRequest, StateUpdate, StopReason, }; use llm::{LlmResponse, testing::FakeLlmProvider}; use std::sync::Arc; @@ -192,33 +193,63 @@ async fn invalid_live_resume_leaves_original_turn_running() { .await; } +#[tokio::test(flavor = "current_thread")] +async fn elicitations_reach_the_client_and_its_answer_reaches_the_server() { + AcpTestHarness::run(|mut harness| async move { + let mut results = harness.elicit_during_prompt_expansion(); + harness.peer.queue_elicitation_response( + serde_json::from_value(serde_json::json!({"action": "accept", "content": {"approved": true}})).unwrap(), + ); + let session = harness.insert_agent_switching_session().await; + let id = session.session_id().clone(); + + harness + .client_cx + .send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])) + .block_task() + .await + .expect("the prompt completes"); + + let request = harness.peer.next_elicitation_request().await; + assert_eq!(acp_utils::elicitation::source_mcp_server_name(request.meta.as_ref()), Some("planner-mcp")); + let ElicitationMode::Form(form) = request.mode else { panic!("expected a form elicitation") }; + let ElicitationScope::Session(scope) = form.scope else { panic!("expected a session-scoped elicitation") }; + assert_eq!(scope.session_id, id); + let answer = results.recv().await.expect("the server receives the answer"); + assert_eq!(answer.action, mcp_utils::model::ElicitationAction::Accept); + assert_eq!(answer.content, Some(serde_json::json!({"approved": true}))); + harness.shutdown().await; + }) + .await; +} + #[tokio::test(flavor = "current_thread")] async fn pending_and_detached_elicitations_cancel_without_stopping_session() { LocalSet::new() .run_until(async { for detach_before_elicitation in [false, true] { let mut harness = AcpTestHarness::start().await; - let (started, release) = harness.pause_prompt_expansion(); + let expansion = harness.pause_prompt_expansion(); let mut results = harness.elicit_during_prompt_expansion(); let session = harness.insert_agent_switching_session().await; let id = session.session_id().clone(); let capture = harness.peer.capture_next_elicitation(); let prompt = harness.client_cx.send_request(PromptRequest::new(id.clone(), vec!["/plan".into()])).block_task(); - started.notified().await; + expansion.reached().await; let pending = if detach_before_elicitation { harness.disconnect().await; - release.notify_one(); + expansion.open(); None } else { - release.notify_one(); + expansion.open(); let responder = capture.await.expect("client received elicitation"); harness.disconnect().await; Some(responder) }; assert_eq!( results.recv().await.expect("MCP receives cancellation"), - mcp_utils::client::cancel_result() + mcp_utils::model::ElicitResult::new(mcp_utils::model::ElicitationAction::Cancel) ); drop(pending); assert!(prompt.await.is_err(), "disconnected prompt response is not retried"); diff --git a/crates/aether-cli/tests/integration/acp_session_lifecycle.rs b/crates/aether-cli/tests/integration/acp_session_lifecycle.rs index ddfbed9dd..632d17dc3 100644 --- a/crates/aether-cli/tests/integration/acp_session_lifecycle.rs +++ b/crates/aether-cli/tests/integration/acp_session_lifecycle.rs @@ -212,6 +212,21 @@ async fn resume_restores_transcript_without_replay_and_replaces_active_session() .await; } +#[tokio::test(flavor = "current_thread")] +async fn commands_are_advertised_while_another_mcp_server_is_still_connecting() { + AcpTestHarness::run(|mut harness| async move { + harness.add_stalled_mcp_server(); + let session = harness.insert_agent_switching_session().await; + + let advertises_plan = |update: &SessionUpdate| { + matches!(update, SessionUpdate::AvailableCommandsUpdate(commands) + if commands.available_commands.iter().any(|command| command.name == "plan")) + }; + harness.updates_until(session.session_id(), advertises_plan).await; + }) + .await; +} + #[tokio::test(flavor = "current_thread")] async fn resume_and_close_reject_unknown_sessions() { AcpTestHarness::run(|harness| async move { diff --git a/crates/aether-cli/tests/integration/acp_slash_commands.rs b/crates/aether-cli/tests/integration/acp_slash_commands.rs new file mode 100644 index 000000000..340fce8b1 --- /dev/null +++ b/crates/aether-cli/tests/integration/acp_slash_commands.rs @@ -0,0 +1,108 @@ +use aether_cli::acp::testing::{AcpTestHarness, FakeAgentSwitchingSession}; +use agent_client_protocol::schema::v2::{AvailableCommand, AvailableCommandInput, PromptRequest, StopReason}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::{Prompt, PromptArgument}; +use mcp_utils::testing::FakeMcpServer; + +#[tokio::test(flavor = "current_thread")] +async fn slash_command_arguments_fill_the_prompt_template() { + AcpTestHarness::run(|mut harness| async move { + harness.add_mcp_server(ServerSpec::in_process( + "docs", + FakeMcpServer::new().with_prompt("greet", "{ARGUMENTS}|{1}|{3}"), + )); + let session = harness.insert_agent_switching_session().await; + + send(&mut harness, &session, "/greet do a thing").await; + + session.planner().assert_saw(&["do a thing|do|thing"]); + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn unknown_slash_commands_reach_the_agent_unchanged() { + AcpTestHarness::run(|mut harness| async move { + let session = harness.insert_agent_switching_session().await; + + send(&mut harness, &session, "/missing").await; + + session.planner().assert_saw(&["/missing"]); + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn commands_keep_the_names_their_servers_gave_them() { + AcpTestHarness::run(|mut harness| async move { + harness.add_mcp_server(ServerSpec::in_process( + "docs", + FakeMcpServer::new().with_prompt("release__notes", "notes for {1}"), + )); + harness.add_mcp_server(ServerSpec::in_process( + "deep__docs", + FakeMcpServer::new().with_prompt("outline", "outline for {1}"), + )); + let session = harness.insert_agent_switching_session().await; + + harness.expect_available_commands(&["release__notes", "outline"], &["notes", "deep__docs__outline"]).await; + send(&mut harness, &session, "/release__notes v2").await; + send(&mut harness, &session, "/outline v3").await; + + session.planner().assert_saw(&["notes for v2", "outline for v3"]); + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn the_first_server_offering_a_prompt_name_owns_the_command() { + AcpTestHarness::run(|mut harness| async move { + harness + .add_mcp_server(ServerSpec::in_process("second", FakeMcpServer::new().with_prompt("plan", "second plan"))); + let session = harness.insert_agent_switching_session().await; + + let commands = harness.next_available_commands().await; + send(&mut harness, &session, "/plan").await; + + assert_eq!(commands.iter().filter(|command| command.name == "plan").count(), 1); + session.planner().assert_saw(&["expanded plan"]); + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn commands_carry_the_prompt_description_and_argument_hint() { + AcpTestHarness::run(|mut harness| async move { + let argument = PromptArgument::new("ARGUMENTS").with_description("[query]").with_required(false); + let search = Prompt::new("search", Some("Search code"), Some(vec![argument])); + harness.add_mcp_server(ServerSpec::in_process( + "code", + FakeMcpServer::new().with_prompt_definition(search, "search"), + )); + harness.insert_agent_switching_session().await; + + let commands = harness.next_available_commands().await; + + assert_eq!(command(&commands, "search").description, "Search code"); + assert_eq!(hint(&commands, "search"), "[query]"); + assert_eq!(hint(&commands, "plan"), "optional arguments"); + }) + .await; +} + +async fn send(harness: &mut AcpTestHarness, session: &FakeAgentSwitchingSession, text: &str) { + let id = session.session_id().clone(); + harness.client_cx.send_request(PromptRequest::new(id.clone(), vec![text.into()])).block_task().await.unwrap(); + harness.expect_idle(&id, StopReason::EndTurn).await; +} + +fn command<'a>(commands: &'a [AvailableCommand], name: &str) -> &'a AvailableCommand { + commands.iter().find(|command| command.name == name).unwrap_or_else(|| panic!("expected /{name} in {commands:?}")) +} + +fn hint(commands: &[AvailableCommand], name: &str) -> String { + match &command(commands, name).input { + Some(AvailableCommandInput::Text(input)) => input.hint.clone(), + other => panic!("expected /{name} to take text input, got {other:?}"), + } +} diff --git a/crates/aether-cli/tests/integration/main.rs b/crates/aether-cli/tests/integration/main.rs index 2d03a3b9f..3925adf1d 100644 --- a/crates/aether-cli/tests/integration/main.rs +++ b/crates/aether-cli/tests/integration/main.rs @@ -2,15 +2,16 @@ mod acp_agent_switching; mod acp_cancellation; mod acp_control_plane; mod acp_git_diff; +mod acp_mcp_sign_in; mod acp_prompt_lifecycle; mod acp_prompt_search; mod acp_remote; mod acp_session_lifecycle; +mod acp_slash_commands; mod acp_stdio; mod acp_workspace_move; mod mcp_command; mod run_init; -mod slash_commands; mod trace_context; type TestResult = Result>; diff --git a/crates/aether-cli/tests/integration/mcp_command.rs b/crates/aether-cli/tests/integration/mcp_command.rs index 53d0eebda..788b2b8d5 100644 --- a/crates/aether-cli/tests/integration/mcp_command.rs +++ b/crates/aether-cli/tests/integration/mcp_command.rs @@ -1,77 +1,22 @@ +use aether_core::mcp::{AETHER_MCP_IPC_SOCKET, McpCatalog, ToolExposure, mcp}; +use aether_core::testing::ready_catalog; use mcp_servers::coding::tools::bash::{BashEnvironment, BashInput, execute_command}; -use mcp_utils::tool_gateway::{AETHER_MCP_IPC_SOCKET, UnixSocketMcpTransport, UnixSocketPath}; -use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ErrorData, Implementation, ListToolsResult, - PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, -}; -use rmcp::{RoleServer, ServerHandler, service::RequestContext}; -use serde_json::{Map, json}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::CallToolResult; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse}; +use serde_json::json; use std::os::unix::fs::PermissionsExt; use std::path::Path; use std::process::{Command, Output, Stdio}; -use std::sync::Arc; - -#[derive(Clone)] -struct FakeGateway; - -impl ServerHandler for FakeGateway { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("fake-gateway", "1.0.0")) - } - - fn list_tools( - &self, - _request: Option, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - let schema = json!({ - "type": "object", - "properties": { "message": { "type": "string", "description": "Text to echo" } }, - "required": ["message"] - }); - std::future::ready(Ok(ListToolsResult::with_all_items(vec![ - Tool::new("_aether_list_servers", "List deferred servers", Arc::new(Map::new())), - Tool::new("demo__echo", "Echo a message", Arc::new(schema.as_object().unwrap().clone())), - Tool::new("demo__fail", "Return a tool error", Arc::new(Map::new())), - Tool::new("demo__slow", "Wait for cancellation", Arc::new(Map::new())), - Tool::new("demo__text", "Return content blocks", Arc::new(Map::new())), - ]))) - } - - async fn call_tool( - &self, - request: CallToolRequestParams, - _context: RequestContext, - ) -> Result { - match request.name.as_ref() { - "_aether_list_servers" => Ok(CallToolResult::structured(json!([ - {"name": "demo", "description": "Demonstration tools"} - ])) - .into()), - "demo__echo" => Ok(CallToolResult::structured(json!({ - "echoed": request.arguments.unwrap_or_default().get("message").cloned().unwrap_or(json!(null)) - })) - .into()), - "demo__fail" => Ok(CallToolResult::error(vec![]).into()), - "demo__slow" => { - std::future::pending::<()>().await; - Ok(CallToolResult::structured(json!({"done": true})).into()) - } - "demo__text" => Ok(CallToolResult::success(vec![ContentBlock::text("plain")]).into()), - _ => Err(ErrorData::invalid_params("unknown tool", None)), - } - } -} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn layered_help_reveals_servers_tools_and_complete_schema() { - let (_server, socket) = gateway(); + let (_session, socket) = gateway().await; let top = run(&socket, &["mcp", "--help"], None); assert_success(&top); assert_stdout_contains(&top, "demo"); - assert_stdout_contains(&top, "Demonstration tools"); + assert_stdout_contains(&top, "A fake MCP server for testing"); let server = run(&socket, &["mcp", "demo", "--help"], None); assert_success(&server); @@ -85,9 +30,39 @@ async fn layered_help_reveals_servers_tools_and_complete_schema() { assert_stdout_contains(&tool, "--json"); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn server_help_supports_namespace_delimiters_in_server_names() { + let (_session, socket) = gateway().await; + + let server = run(&socket, &["mcp", "deep__demo", "--help"], None); + assert_success(&server); + assert_stdout_contains(&server, " ping"); + assert_stdout_contains(&server, "Ping a nested server"); + assert!(!String::from_utf8_lossy(&server.stdout).contains("demo__ping")); + + let demo = run(&socket, &["mcp", "demo", "--help"], None); + assert_success(&demo); + assert!(!String::from_utf8_lossy(&demo.stdout).contains("ping")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn server_help_lists_only_the_named_servers_tools() { + let (_session, socket) = gateway().await; + + let docs = run(&socket, &["mcp", "docs", "--help"], None); + assert_success(&docs); + assert_stdout_contains(&docs, "Search the docs"); + assert!(!String::from_utf8_lossy(&docs.stdout).contains("Search the v2 docs")); + + let v2 = run(&socket, &["mcp", "docs__v2", "--help"], None); + assert_success(&v2); + assert_stdout_contains(&v2, "Search the v2 docs"); + assert!(!String::from_utf8_lossy(&v2.stdout).contains("Search the docs")); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn calls_accept_empty_json_flag_or_stdin_and_print_one_json_value() { - let (_server, socket) = gateway(); + let (_session, socket) = gateway().await; let empty = run(&socket, &["mcp", "demo", "echo"], None); assert_success(&empty); @@ -108,7 +83,7 @@ async fn calls_accept_empty_json_flag_or_stdin_and_print_one_json_value() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn json_flag_does_not_wait_for_open_stdin() { - let (_server, socket) = gateway(); + let (_session, socket) = gateway().await; let mut command = Command::new(env!("CARGO_BIN_EXE_aether")); command .args(["mcp", "demo", "echo", "--json", r#"{"message":"flag"}"#]) @@ -126,7 +101,7 @@ async fn json_flag_does_not_wait_for_open_stdin() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn coding_bash_composes_real_cli_with_jq_pipes_redirects_and_scripts() { - let (_server, socket) = gateway(); + let (_session, socket) = gateway().await; let workspace = tempfile::tempdir().unwrap(); let script = workspace.path().join("call.sh"); std::fs::write( @@ -183,7 +158,7 @@ async fn coding_bash_composes_real_cli_with_jq_pipes_redirects_and_scripts() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn usage_errors_exit_two_and_runtime_errors_exit_one() { - let (_server, socket) = gateway(); + let (_session, socket) = gateway().await; for args in [ vec!["mcp", "demo"], @@ -222,11 +197,37 @@ async fn usage_errors_exit_two_and_runtime_errors_exit_one() { assert!(String::from_utf8_lossy(&missing.stderr).contains(AETHER_MCP_IPC_SOCKET)); } -fn gateway() -> (mcp_utils::tool_gateway::UnixSocketServer, String) { - let path = UnixSocketPath::new().unwrap(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - let socket = transport.path().to_string_lossy().into_owned(); - (transport.spawn(FakeGateway), socket) +/// A session whose servers all defer their tools, as `aether mcp` reaches them through the session's socket. +async fn gateway() -> (McpCatalog, String) { + let demo = FakeMcpServer::new() + .with_tool( + FakeTool::new("echo") + .description("Echo a message") + .input_schema(json!({ + "type": "object", + "properties": { "message": { "type": "string", "description": "Text to echo" } }, + "required": ["message"] + })) + .responds_with(|request| { + let message = request.arguments.as_ref().and_then(|args| args.get("message")).cloned(); + CallToolResult::structured(json!({ "echoed": message })).into() + }), + ) + .with_tool(FakeTool::new("fail").responds(CallToolResult::error(vec![]))) + .with_tool(FakeTool::new("slow").responds(FakeToolResponse::pending())) + .with_tool(FakeTool::new("text").responds(FakeToolResponse::text("plain"))); + let servers = [ + ("demo", demo), + ("deep__demo", FakeMcpServer::new().with_tool(FakeTool::new("ping").description("Ping a nested server"))), + ("docs", FakeMcpServer::new().with_tool(FakeTool::new("search").description("Search the docs"))), + ("docs__v2", FakeMcpServer::new().with_tool(FakeTool::new("search").description("Search the v2 docs"))), + ]; + let builder = servers.into_iter().fold(mcp("/tmp"), |builder, (name, server)| { + builder.with_server(ServerSpec::in_process(name, server), ToolExposure::deferred_all()) + }); + let catalog = ready_catalog(builder).await; + let socket = catalog.deferred_tools_socket().expect("deferred servers are served").to_string_lossy().into_owned(); + (catalog, socket) } fn run(socket: &str, args: &[&str], stdin: Option<&str>) -> Output { diff --git a/crates/aether-cli/tests/integration/run_init.rs b/crates/aether-cli/tests/integration/run_init.rs index b321bced9..3bc0717e4 100644 --- a/crates/aether-cli/tests/integration/run_init.rs +++ b/crates/aether-cli/tests/integration/run_init.rs @@ -1,11 +1,10 @@ use aether_cli::init::{HarnessIntegration, InitError, InitOutcome, InitTarget, OverwriteMode, Preset, apply_init}; use aether_core::core::Prompt; +use aether_core::mcp::{McpConfig, McpServerConfig, ToolMatcher, mcp}; use aether_project::{AetherSettings, AgentCatalog, McpSourceSpec, PromptSource}; use llm::catalog::Provider; use llm::{LlmModel, ReasoningEffort}; -use mcp_servers::{CodingMcpArgs, ReviewMcpArgs, SkillsMcpArgs, SubAgentsMcpArgs, TasksMcpArgs}; -use mcp_utils::client::McpServerConfig; -use mcp_utils::client::ToolMatcher; +use mcp_servers::McpBuilderExt; use std::collections::BTreeMap; use std::path::Path; @@ -26,10 +25,10 @@ fn inline_args(servers: &BTreeMap, name: &str) -> Vec &BTreeMap { - let McpSourceSpec::Inline { servers } = spec else { + let McpSourceSpec::Inline(config) = spec else { panic!("expected inline MCPs, got a file spec"); }; - servers + &config.servers } #[test] @@ -277,8 +276,8 @@ fn force_overwrites_selected_target_only() { assert_eq!(std::fs::read_to_string(dir.path().join(".aether/settings.json")).unwrap(), "{}"); } -#[test] -fn every_inline_mcp_in_init_presets_parses_its_args() { +#[tokio::test] +async fn every_inline_mcp_in_init_presets_spawns_with_its_args() { for preset in [Preset::Minimal, Preset::BatteriesIncluded] { for harnesses in [vec![], vec![HarnessIntegration::Claude], vec![HarnessIntegration::Claude, HarnessIntegration::Agents]] @@ -295,28 +294,16 @@ fn every_inline_mcp_in_init_presets_parses_its_args() { let settings = load(&dir.path().join("settings.json")); for agent in &settings.agents { - for mcp in &agent.mcps { - let McpSourceSpec::Inline { servers } = mcp else { continue }; - for (name, config) in servers { - let McpServerConfig::InMemory(in_memory) = config else { continue }; - let args = in_memory.args.clone(); - let parse_result = match name.as_str() { - "coding" => CodingMcpArgs::from_args(args).map(|_| ()), - "skills" => SkillsMcpArgs::from_args(args).map(|_| ()), - "subagents" => SubAgentsMcpArgs::from_args(args).map(|_| ()), - "tasks" => TasksMcpArgs::from_args(args).map(|_| ()), - "review" => ReviewMcpArgs::from_args(args).map(|_| ()), - other => { - panic!("preset references unknown in-memory MCP `{other}`; add a parse check") - } - }; - assert!( - parse_result.is_ok(), - "{preset:?} agent `{}` server `{name}` harnesses={harnesses:?} args failed to parse: {parse_result:?}", - agent.name - ); - } - } + let inline = agent.mcps.iter().filter_map(|source| match source { + McpSourceSpec::Inline(config) => Some(config.clone()), + McpSourceSpec::File(_) => None, + }); + let config = inline.collect::(); + let spawned = mcp(dir.path()).with_builtin_servers().with_config(config).unwrap().spawn(); + let catalog = spawned.unwrap_or_else(|error| { + panic!("{preset:?} agent `{}` harnesses={harnesses:?} failed to spawn: {error}", agent.name) + }); + catalog.shutdown().await; } } } @@ -369,8 +356,8 @@ fn assert_read_only_coding_tools(agent: &aether_project::AgentConfig) { } fn assert_minimal_mcp_and_tools(agent: &aether_project::AgentConfig) { - let McpSourceSpec::Inline { servers } = &agent.mcps[0] else { panic!("minimal MCPs should be inline") }; - let names: Vec<&str> = servers.keys().map(String::as_str).collect(); + let McpSourceSpec::Inline(config) = &agent.mcps[0] else { panic!("minimal MCPs should be inline") }; + let names: Vec<&str> = config.servers.keys().map(String::as_str).collect(); assert_eq!(names, vec!["coding", "skills"]); assert_eq!(agent.tools.allow, vec![ToolMatcher::name("coding__bash"), ToolMatcher::name("skills__*")]); assert_eq!(agent.tools.deny.len(), 0); diff --git a/crates/aether-cli/tests/integration/slash_commands.rs b/crates/aether-cli/tests/integration/slash_commands.rs deleted file mode 100644 index 2c5807105..000000000 --- a/crates/aether-cli/tests/integration/slash_commands.rs +++ /dev/null @@ -1,49 +0,0 @@ -use agent_client_protocol::schema::v2 as acp; -use rmcp::model::{Prompt, PromptArgument}; - -/// Create an MCP prompt with no arguments. -fn create_prompt(name: &str, description: Option<&str>) -> Prompt { - Prompt::new(name.to_string(), description.map(str::to_string), None) -} - -/// Create an MCP prompt with an ARGUMENTS parameter (unified prompt format). -fn create_prompt_with_hint(name: &str, description: Option<&str>, hint: &str) -> Prompt { - let arguments = Some(vec![PromptArgument::new("ARGUMENTS").with_description(hint).with_required(false)]); - Prompt::new(name.to_string(), description.map(str::to_string), arguments) -} - -#[test] -fn test_map_prompt_to_command_strips_namespace() { - let prompt = create_prompt("coding__web", Some("Search the web")); - - let command = aether_cli::map_mcp_prompt_to_available_command(&prompt); - - assert_eq!(command.name, "web"); - assert_eq!(command.description, "Search the web"); - assert!(command.input.is_some()); -} - -#[test] -fn test_map_prompt_to_command_without_namespace() { - let prompt = create_prompt("test", Some("Run tests")); - - let command = aether_cli::map_mcp_prompt_to_available_command(&prompt); - - assert_eq!(command.name, "test"); - assert_eq!(command.description, "Run tests"); -} - -#[test] -fn test_map_prompt_to_command_with_argument_hint() { - let prompt = create_prompt_with_hint("search", Some("Search code"), "[query]"); - - let command = aether_cli::map_mcp_prompt_to_available_command(&prompt); - - assert_eq!(command.name, "search"); - match command.input { - Some(acp::AvailableCommandInput::Text(input)) => { - assert_eq!(input.hint, "[query]"); - } - _ => panic!("Expected Text input with hint"), - } -} diff --git a/crates/aether-core/Cargo.toml b/crates/aether-core/Cargo.toml index 5ec06394e..eba968025 100644 --- a/crates/aether-core/Cargo.toml +++ b/crates/aether-core/Cargo.toml @@ -26,6 +26,7 @@ path = "src/lib.rs" [dependencies] async-stream = { workspace = true } futures = { workspace = true } +tokio-util = { workspace = true, features = ["rt"] } serde = { workspace = true } serde_json = { workspace = true } noyalib = { workspace = true } @@ -38,6 +39,7 @@ schemars = { workspace = true } tracing = { workspace = true } thiserror = { workspace = true } glob = { workspace = true } +http = { workspace = true } aether-auth = { path = "../aether-auth", version = "0.1.17" } mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", version = "0.7.13" } llm = { package = "aether-llm", path = "../llm", version = "0.9.13", features = ["codex"] } diff --git a/crates/aether-core/README.md b/crates/aether-core/README.md index 8efeba6c5..857a8ea74 100644 --- a/crates/aether-core/README.md +++ b/crates/aether-core/README.md @@ -106,6 +106,7 @@ use aether_core::core::{Command, Prompt, agent}; use aether_core::events::{AgentEvent, MessageEvent, ToolEvent, TurnEvent}; use aether_core::mcp::mcp; use llm::providers::openrouter::OpenRouterProvider; +use aether_core::mcp::McpConfig; use std::io::{self, Write}; #[tokio::main] @@ -113,16 +114,15 @@ async fn main() -> Result<(), Box> { let llm = OpenRouterProvider::default("z-ai/glm-4.6")?; // 1. Connect to MCP servers - let mut mcp_runtime = mcp(".") - .from_json_files(&["mcp.json"])? // <-- Load MCP servers from one or more JSON files - .spawn() // <-- Spawn the MCP client into a tokio task (multiple agents can use it) - .await?; - let snapshot = mcp_runtime.block_until_ready().await.expect("MCP should initialize"); + let mcp = mcp(".") + .with_config(McpConfig::from_json_file("mcp.json")?)? // <-- Load MCP servers from a JSON file + .spawn()?; // <-- Start the MCP gateway; servers connect in the background + mcp.ready().await; // <-- Wait for every server to connect // 2. Create Agent let (tx, mut rx, _handle) = agent(llm) .system_prompt(Prompt::file("AGENTS.md", ".")) // <-- Load system prompt from AGENTS.md - .tools(mcp_runtime.handle().clone(), snapshot.tool_definitions()) // <-- Give the agent MCP tools + .mcp(mcp) // <-- Give the agent MCP tools, kept in sync as they change .spawn() .await?; diff --git a/crates/aether-core/examples/mcp_agent.rs b/crates/aether-core/examples/mcp_agent.rs index 4c361ad94..fa58edcfe 100644 --- a/crates/aether-core/examples/mcp_agent.rs +++ b/crates/aether-core/examples/mcp_agent.rs @@ -1,4 +1,5 @@ use aether_core::events::{AgentEvent, ContextEvent, MessageEvent, ModelEvent, ToolEvent, TurnEvent}; +use aether_core::mcp::McpConfig; use aether_core::{ core::{Prompt, agent}, events::{Command, TurnOutcome, UserCommand}, @@ -14,12 +15,13 @@ async fn main() -> Result<(), Box> { tracing_subscriber::fmt::init(); let llm = OpenRouterProvider::default("z-ai/glm-4.5-air")?; - let mut spawn = mcp(std::env::current_dir()?).from_json_files(&["examples/mcp.json"])?.spawn().await?; - let connection_details = spawn.block_until_ready().await.ok_or("MCP bootstrap aborted before completion")?; + let config = McpConfig::from_json_file("examples/mcp.json")?; + let mcp = mcp(std::env::current_dir()?).with_config(config)?.spawn()?; + mcp.ready().await; let (tx, mut rx, _handle) = agent(llm) .system_prompt(Prompt::text("You are a helpful assistant with access to web browsing tools via Playwright.")) - .tools(spawn.handle().clone(), connection_details.tool_definitions()) + .mcp(mcp) .spawn() .await?; diff --git a/crates/aether-core/src/agent_spec.rs b/crates/aether-core/src/agent_spec.rs index dae2eb85d..46416c91b 100644 --- a/crates/aether-core/src/agent_spec.rs +++ b/crates/aether-core/src/agent_spec.rs @@ -4,30 +4,8 @@ //! It represents a resolved runtime type, not a raw settings DTO. use crate::core::Prompt; +use crate::mcp::{McpConfig, ToolFilter}; use llm::{LlmModel, ModelSettings, ProviderConnectionOverrides, ReasoningEffort}; -use mcp_utils::client::{McpConfig, ToolFilter}; -use std::path::PathBuf; - -#[derive(Debug, Clone)] -pub enum McpConfigSource { - File { path: PathBuf, defer_tools: bool }, - Json(String), - Inline(McpConfig), -} - -impl McpConfigSource { - pub fn file(path: PathBuf, defer_tools: bool) -> Self { - Self::File { path, defer_tools } - } - - pub fn model_visible(path: PathBuf) -> Self { - Self::file(path, false) - } - - pub fn deferred(path: PathBuf) -> Self { - Self::file(path, true) - } -} /// A resolved agent specification ready for runtime use. /// @@ -55,11 +33,8 @@ pub struct AgentSpec { pub prompts: Vec, /// Provider connection overrides keyed by model provider name. pub provider_connections: ProviderConnectionOverrides, - /// Resolved MCP config sources for this agent, applied in order. - /// - /// Model-visible server name collisions use last-source-wins semantics. File sources - /// can defer all of their tools for progressive discovery. - pub mcp_config_sources: Vec, + /// MCP servers this agent connects to + pub mcp: McpConfig, /// How this agent can be invoked. pub exposure: AgentSpecExposure, /// Tool filter for restricting which MCP tools this agent can use. @@ -79,7 +54,7 @@ impl AgentSpec { context_window: None, prompts, provider_connections: ProviderConnectionOverrides::default(), - mcp_config_sources: Vec::new(), + mcp: McpConfig::default(), exposure: AgentSpecExposure::none(), tools: ToolFilter::default(), } @@ -124,8 +99,7 @@ impl AgentSpecExposure { #[cfg(test)] mod tests { use super::*; - use llm::{ToolAnnotations, ToolDefinition}; - use mcp_utils::client::{ToolAnnotationMatcher, ToolMatcher}; + use std::path::PathBuf; #[test] fn default_spec_has_expected_fields() { @@ -138,201 +112,7 @@ mod tests { assert_eq!(spec.model, model.to_string()); assert!(spec.reasoning_effort.is_none()); assert_eq!(spec.prompts.len(), 1); - assert!(spec.mcp_config_sources.is_empty()); + assert!(spec.mcp.is_empty()); assert_eq!(spec.exposure, AgentSpecExposure::none()); } - - fn make_tool(name: &str) -> ToolDefinition { - ToolDefinition::new(name, "", serde_json::json!({})) - } - - fn make_annotated_tool(name: &str, annotations: ToolAnnotations) -> ToolDefinition { - ToolDefinition::new(name, "", serde_json::json!({})).with_annotations(annotations) - } - - #[test] - fn empty_filter_allows_all_tools() { - let filter = ToolFilter::default(); - let tools = vec![make_tool("bash"), make_tool("read_file")]; - let result = filter.apply(tools); - assert_eq!(result.len(), 2); - } - - #[test] - fn allow_keeps_only_matching_tools() { - let filter = - ToolFilter { allow: vec![ToolMatcher::name("read_file"), ToolMatcher::name("grep")], deny: vec![] }; - let tools = vec![make_tool("bash"), make_tool("read_file"), make_tool("grep")]; - let result = filter.apply(tools); - let names: Vec<_> = result.iter().map(|t| t.name.as_str()).collect(); - assert_eq!(names, vec!["read_file", "grep"]); - } - - #[test] - fn deny_removes_matching_tools() { - let filter = ToolFilter { allow: vec![], deny: vec![ToolMatcher::name("bash")] }; - let tools = vec![make_tool("bash"), make_tool("read_file")]; - let result = filter.apply(tools); - let names: Vec<_> = result.iter().map(|t| t.name.as_str()).collect(); - assert_eq!(names, vec!["read_file"]); - } - - #[test] - fn wildcard_matching() { - let filter = ToolFilter { allow: vec![ToolMatcher::name("coding__*")], deny: vec![] }; - let tools = vec![make_tool("coding__grep"), make_tool("coding__read_file"), make_tool("plugins__bash")]; - let result = filter.apply(tools); - let names: Vec<_> = result.iter().map(|t| t.name.as_str()).collect(); - assert_eq!(names, vec!["coding__grep", "coding__read_file"]); - } - - #[test] - fn combined_allow_and_deny() { - let filter = ToolFilter { - allow: vec![ToolMatcher::name("coding__*")], - deny: vec![ToolMatcher::name("coding__write_file")], - }; - let tools = vec![ - make_tool("coding__grep"), - make_tool("coding__write_file"), - make_tool("coding__read_file"), - make_tool("plugins__bash"), - ]; - let result = filter.apply(tools); - let names: Vec<_> = result.iter().map(|t| t.name.as_str()).collect(); - assert_eq!(names, vec!["coding__grep", "coding__read_file"]); - } - - #[test] - fn annotation_allow_matches_present_values() { - let filter = ToolFilter { allow: vec![ToolMatcher::read_only()], deny: vec![] }; - let tools = vec![ - make_tool("unknown"), - make_annotated_tool("read", ToolAnnotations { read_only_hint: Some(true), ..ToolAnnotations::default() }), - make_annotated_tool("write", ToolAnnotations { read_only_hint: Some(false), ..ToolAnnotations::default() }), - ]; - let names: Vec<_> = filter.apply(tools).into_iter().map(|tool| tool.name).collect(); - assert_eq!(names, vec!["read"]); - } - - #[test] - fn deny_annotation_removes_destructive_tools() { - let filter = ToolFilter { - allow: vec![], - deny: vec![ToolMatcher::annotations(ToolAnnotationMatcher { - destructive: Some(true), - ..ToolAnnotationMatcher::default() - })], - }; - let tools = vec![ - make_tool("unknown"), - make_annotated_tool( - "safe_update", - ToolAnnotations { - read_only_hint: Some(false), - destructive_hint: Some(false), - ..ToolAnnotations::default() - }, - ), - ]; - let names: Vec<_> = filter.apply(tools).into_iter().map(|tool| tool.name).collect(); - assert_eq!(names, vec!["unknown", "safe_update"]); - } - - #[test] - fn annotation_matchers_do_not_match_missing_fields() { - let filter = ToolFilter { - allow: vec![], - deny: vec![ - ToolMatcher::annotations(ToolAnnotationMatcher { - destructive: Some(true), - ..ToolAnnotationMatcher::default() - }), - ToolMatcher::annotations(ToolAnnotationMatcher { - open_world: Some(true), - ..ToolAnnotationMatcher::default() - }), - ToolMatcher::annotations(ToolAnnotationMatcher { - idempotent: Some(false), - ..ToolAnnotationMatcher::default() - }), - ToolMatcher::annotations(ToolAnnotationMatcher { - read_only: Some(false), - ..ToolAnnotationMatcher::default() - }), - ], - }; - let tools = vec![make_tool("unknown")]; - let names: Vec<_> = filter.apply(tools).into_iter().map(|tool| tool.name).collect(); - assert_eq!(names, vec!["unknown"]); - } - - #[test] - fn annotation_matchers_do_not_infer_fields_from_read_only_hint() { - let filter = ToolFilter { - allow: vec![ToolMatcher::annotations(ToolAnnotationMatcher { - destructive: Some(false), - ..ToolAnnotationMatcher::default() - })], - deny: vec![], - }; - let tools = vec![make_annotated_tool("read", ToolAnnotations::read_only())]; - assert_eq!(filter.apply(tools).len(), 0); - } - - #[test] - fn deny_wins_over_allow() { - let filter = - ToolFilter { allow: vec![ToolMatcher::read_only()], deny: vec![ToolMatcher::name("coding__read_file")] }; - let tools = vec![make_annotated_tool( - "coding__read_file", - ToolAnnotations { read_only_hint: Some(true), ..ToolAnnotations::default() }, - )]; - assert_eq!(filter.apply(tools).len(), 0); - } - - #[test] - fn mixed_allow_entries_are_ored() { - let filter = ToolFilter { allow: vec![ToolMatcher::read_only(), ToolMatcher::name("review__*")], deny: vec![] }; - let tools = vec![ - make_annotated_tool( - "coding__grep", - ToolAnnotations { read_only_hint: Some(true), ..ToolAnnotations::default() }, - ), - make_tool("review__review_artifact"), - make_tool("coding__bash"), - ]; - let names: Vec<_> = filter.apply(tools).into_iter().map(|tool| tool.name).collect(); - assert_eq!(names, vec!["coding__grep", "review__review_artifact"]); - } - - #[test] - fn empty_annotation_matcher_matches_nothing() { - let filter = - ToolFilter { allow: vec![ToolMatcher::annotations(ToolAnnotationMatcher::default())], deny: vec![] }; - let tools = vec![make_annotated_tool( - "coding__grep", - ToolAnnotations { read_only_hint: Some(true), ..ToolAnnotations::default() }, - )]; - assert_eq!(filter.apply(tools).len(), 0); - } - - #[test] - fn exact_name_match_is_not_a_prefix_match() { - let filter = ToolFilter { allow: vec![ToolMatcher::name("bash")], deny: vec![] }; - let names: Vec<_> = - filter.apply(vec![make_tool("bash"), make_tool("bash_extended")]).into_iter().map(|t| t.name).collect(); - assert_eq!(names, vec!["bash"]); - } - - #[test] - fn tool_matcher_uses_exact_and_trailing_wildcard_names() { - let exact = ToolMatcher::name("foo"); - let wildcard = ToolMatcher::name("foo*"); - assert!(exact.matches(&make_tool("foo"))); - assert!(!exact.matches(&make_tool("foobar"))); - assert!(wildcard.matches(&make_tool("foobar"))); - assert!(wildcard.matches(&make_tool("foo"))); - assert!(!wildcard.matches(&make_tool("bar"))); - } } diff --git a/crates/aether-core/src/core/agent.rs b/crates/aether-core/src/core/agent.rs index 7fccfd027..34aecc639 100644 --- a/crates/aether-core/src/core/agent.rs +++ b/crates/aether-core/src/core/agent.rs @@ -8,16 +8,18 @@ pub use crate::core::retry_config::RetryConfig; use crate::core::tool_execution::{ToolAbortPolicy, ToolExecutionUpdate, ToolExecutions}; use crate::events::{ AgentCommand, AgentEvent, AgentObserver, Command, CompactionId, CompactionOutcome, ContextEvent, LlmCallOutcome, - ModelEvent, StreamState, TaskOutcome, ToolEvent, TraceContext, TurnEvent, TurnOutcome, UserCommand, + ModelEvent, StreamState, TaskOutcome, TaskOutcomeState, ToolEvent, TraceContext, TurnEvent, TurnOutcome, + UserCommand, }; -use crate::mcp::{McpHandle, ToolCallStream}; -use futures::{FutureExt, Stream, future}; +use crate::mcp::McpCatalog; +use crate::mcp::tool_bridge::call_tool; +use futures::{FutureExt, Stream}; use llm::{ AssistantReasoning, ChatMessage, Context, EncryptedReasoningContent, LlmCallPurpose, LlmError, LlmModel, LlmResponse, MessageId, ModelIdentity, StopReason, StreamingModelProvider, TokenUsage, ToolCallError, ToolCallRequest, ToolCallResult, }; -use mcp_utils::client::{CallToolError, CallToolOptions, ToolCallEvent}; +use mcp_utils::client::{ToolCallEvent, ToolCallOptions}; use std::collections::VecDeque; use std::panic::{AssertUnwindSafe, catch_unwind}; use std::pin::Pin; @@ -39,6 +41,7 @@ enum StreamEvent { ToolExecution(ToolCallEvent), Command(Command), InputClosed, + McpChanged, Compaction(Result), } @@ -49,6 +52,7 @@ type EventStream = Pin + Send>>; #[derive(Debug, Clone, PartialEq, Eq, Hash)] enum StreamKey { Input, + Mcp, Llm, Compaction, Tool(String), @@ -57,7 +61,7 @@ enum StreamKey { pub(crate) struct AgentConfig { pub llm: Arc, pub context: Context, - pub mcp: Option, + pub mcp: McpCatalog, pub tool_timeout: Duration, pub compaction_config: Option, pub auto_continue: AutoContinue, @@ -71,7 +75,7 @@ pub(crate) struct AgentConfig { pub struct Agent { llm: Arc, context: Context, - mcp: Option, + mcp: McpCatalog, message_tx: mpsc::Sender, observers: Observers, streams: StreamMap, @@ -104,6 +108,8 @@ impl Agent { .chain(futures::stream::once(async { StreamEvent::InputClosed })); streams.insert(StreamKey::Input, Box::pin(input_stream)); + streams.insert(StreamKey::Mcp, Box::pin(config.mcp.changes().map(|()| StreamEvent::McpChanged))); + let context_limit = config.context_window.or_else(|| config.llm.context_window()); Self { @@ -162,7 +168,7 @@ impl Agent { async fn event_loop(&mut self) { let mut state = IterationState::default(); let mut input_closed = false; - self.emit_tool_definitions().await; + self.sync_mcp().await; while let Some((stream_key, event)) = self.streams.next().await { match event { @@ -182,13 +188,8 @@ impl Agent { self.on_switch_model(new_provider).await; } - StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateTools(tools))) => { - self.context.set_tools(tools); - self.emit_tool_definitions().await; - } - - StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateMcpInstructions { server, body })) => { - self.on_update_instruction(server, body).await; + StreamEvent::McpChanged => { + self.sync_mcp().await; } StreamEvent::Command(Command::AgentCommand(AgentCommand::SetReasoningEffort(effort))) => { @@ -350,11 +351,19 @@ impl Agent { } } - async fn on_update_instruction(&mut self, server: String, body: Option) { - self.prompt_cache.update_mcp_instruction(server, body); - match self.prompt_cache.render().await { - Ok(content) => self.context.set_system_content(content), - Err(e) => tracing::warn!("Failed to rebuild system prompt after instructions update: {e}"), + async fn sync_mcp(&mut self) { + let (instructions, tools) = (self.mcp.instructions(), self.mcp.tool_definitions()); + + if self.prompt_cache.set_mcp_instructions(instructions) { + match self.prompt_cache.render().await { + Ok(content) => self.context.set_system_content(content), + Err(e) => tracing::warn!("Failed to rebuild system prompt after instructions update: {e}"), + } + } + + if tools != *self.context.tools() { + self.context.set_tools(tools.clone()); + self.emit(AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools })).await; } } @@ -518,7 +527,7 @@ impl Agent { } ToolRequestComplete { tool_call } => { - self.handle_tool_completion(tool_call).await; + self.handle_tool_completion(tool_call, state).await; } Done { stop_reason } => { @@ -550,32 +559,28 @@ impl Agent { } } - async fn handle_tool_completion(&mut self, tool_call: ToolCallRequest) { + async fn handle_tool_completion(&mut self, tool_call: ToolCallRequest, state: &mut IterationState) { let cancel = self.tool_executions.start(tool_call.clone()); let tool_id = tool_call.id.clone(); tracing::debug!("Tool execution started: {} ({})", tool_call.name, tool_id); self.emit(AgentEvent::Tool(ToolEvent::Call { request: tool_call.clone() })).await; - let events: ToolCallStream = match self.mcp.clone() { - Some(mcp) => { - let trace_context = self.observers.tool_trace_context(&tool_id); - let options = CallToolOptions { - timeout: self.tool_timeout, - meta: trace_context.as_ref().map(TraceContext::to_meta), - cancel, - }; - - mcp.call_model_visible(tool_call.name, &tool_call.arguments, options) - } - - None => Box::pin(futures::stream::once(future::ready(ToolCallEvent::Complete(Err( - CallToolError::Unavailable { message: "MCP runtime is not available".to_string() }, - ))))), + let trace_context = self.observers.tool_trace_context(&tool_id); + let options = ToolCallOptions { + timeout: Some(self.tool_timeout), + meta: trace_context.as_ref().map(TraceContext::to_meta), + cancel, }; - - let stream = events.map(StreamEvent::ToolExecution); - self.streams.insert(StreamKey::Tool(tool_id), Box::pin(stream)); + match call_tool(&self.mcp, &tool_call, options) { + Ok(call) => { + self.streams.insert(StreamKey::Tool(tool_id), Box::pin(call.map(StreamEvent::ToolExecution))); + } + Err(rejection) => { + let update = self.tool_executions.reject(&tool_id, &rejection); + self.apply_tool_update(tool_id, update, state).await; + } + } } async fn handle_llm_usage(&mut self, sample: TokenUsage, state: &mut IterationState) { @@ -670,7 +675,12 @@ impl Agent { } async fn on_tool_execution_event(&mut self, tool_id: String, event: ToolCallEvent, state: &mut IterationState) { - match self.tool_executions.on_event(&tool_id, event) { + let update = self.tool_executions.on_event(&tool_id, event); + self.apply_tool_update(tool_id, update, state).await; + } + + async fn apply_tool_update(&mut self, tool_id: String, update: ToolExecutionUpdate, state: &mut IterationState) { + match update { ToolExecutionUpdate::Event(event) => { if let ToolEvent::SubAgentProgress { payload, .. } = &event && let AgentEvent::SessionUsage(child) = &payload.event @@ -689,13 +699,13 @@ impl Agent { state.completed_tool_calls.push(Ok(result)); self.emit(AgentEvent::Tool(event)).await; } - ToolExecutionUpdate::TaskCompleted(outcome) => { - self.streams.remove(&StreamKey::Tool(tool_id)); - self.receive_input(QueuedInput::TaskOutcome(Box::new(outcome)), state).await; - } - ToolExecutionUpdate::TaskCancelled(outcome) => { + ToolExecutionUpdate::TaskFinished(outcome) => { self.streams.remove(&StreamKey::Tool(tool_id)); - self.record_task_outcome(outcome).await; + if matches!(outcome.state, TaskOutcomeState::Cancelled) { + self.record_task_outcome(outcome).await; + } else { + self.receive_input(QueuedInput::TaskOutcome(Box::new(outcome)), state).await; + } } ToolExecutionUpdate::Retired => { self.streams.remove(&StreamKey::Tool(tool_id)); @@ -745,13 +755,6 @@ impl Agent { } } - async fn emit_tool_definitions(&mut self) { - let tools = self.context.tools().clone(); - if !tools.is_empty() { - self.emit(AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools })).await; - } - } - async fn emit(&mut self, message: AgentEvent) { self.observers.on_event(&message); diff --git a/crates/aether-core/src/core/agent_builder.rs b/crates/aether-core/src/core/agent_builder.rs index bb40c7f01..afd5a7f0a 100644 --- a/crates/aether-core/src/core/agent_builder.rs +++ b/crates/aether-core/src/core/agent_builder.rs @@ -3,17 +3,17 @@ use crate::agent_spec::AgentSpec; use crate::context::{CompactionConfig, SessionUsageTracker}; use crate::core::{Agent, AgentDeps, Prompt, PromptCache, Result}; use crate::events::{AgentEvent, AgentObserver, Command, TurnOutcome}; -use crate::mcp::McpHandle; +use crate::mcp::McpCatalog; use llm::parser::ModelProviderParser; -use llm::{ChatMessage, Context, ModelSettings, SessionUsageEvent, StreamingModelProvider, ToolDefinition}; +use llm::{ChatMessage, Context, ModelSettings, SessionUsageEvent, StreamingModelProvider}; use std::sync::Arc; use std::time::Duration; use tokio::sync::mpsc::{self, Receiver, Sender}; -use tokio::task::JoinHandle; +use tokio_util::task::AbortOnDropHandle; -/// Handle for communicating with a running Agent +/// Handle to a running agent's task; dropping it aborts the agent. pub struct AgentHandle { - handle: JoinHandle<()>, + handle: AbortOnDropHandle<()>, } impl AgentHandle { @@ -42,9 +42,8 @@ pub async fn recv_agent_event(rx: &mut Receiver) -> AgentEvent { pub struct AgentBuilder { llm: Arc, prompts: Vec, - tool_definitions: Vec, initial_messages: Vec, - mcp: Option, + mcp: McpCatalog, channel_capacity: usize, tool_timeout: Duration, compaction_config: Option, @@ -62,9 +61,8 @@ impl AgentBuilder { Self { llm, prompts: Vec::new(), - tool_definitions: Vec::new(), initial_messages: Vec::new(), - mcp: None, + mcp: McpCatalog::default(), channel_capacity: 1000, tool_timeout: Duration::from_mins(60), compaction_config: Some(CompactionConfig::default()), @@ -81,15 +79,18 @@ impl AgentBuilder { /// Create a builder from a resolved `AgentSpec`. /// /// The LLM provider is derived from `spec.model` via `ModelProviderParser`. - /// `base_prompts` are prepended before the spec's own prompts. - pub async fn from_spec(spec: &AgentSpec, base_prompts: Vec, deps: &AgentDeps) -> Result { + pub async fn from_spec(spec: &AgentSpec, deps: &AgentDeps) -> Result { let parser = ModelProviderParser::default().with_provider_connections(spec.provider_connections.clone()); let parser = match deps.oauth_credential_store.clone() { Some(store) => parser.with_codex_provider(store), None => parser, }; let (provider, _) = parser.parse(&spec.model).await?; - let mut builder = Self::new(Arc::from(provider)) + Ok(Self::from_spec_with_llm(spec, deps, Arc::from(provider))) + } + + pub fn from_spec_with_llm(spec: &AgentSpec, deps: &AgentDeps, llm: Arc) -> Self { + let mut builder = Self::new(llm) .context_window(spec.context_window) .model_settings(spec.model_settings.clone()) .session_usage(SessionUsageTracker::new(&spec.name)); @@ -101,15 +102,11 @@ impl AgentBuilder { builder = builder.observer(observer); } - for prompt in base_prompts { - builder = builder.system_prompt(prompt); - } - for prompt in &spec.prompts { builder = builder.system_prompt(prompt.clone()); } - Ok(builder) + builder } /// Add a prompt to the system prompt. @@ -120,9 +117,8 @@ impl AgentBuilder { self } - pub fn tools(mut self, mcp: McpHandle, tools: Vec) -> Self { - self.tool_definitions = tools; - self.mcp = Some(mcp); + pub fn mcp(mut self, mcp: McpCatalog) -> Self { + self.mcp = mcp; self } @@ -257,7 +253,7 @@ impl AgentBuilder { messages.extend(self.initial_messages); let (command_tx, command_rx) = mpsc::channel::(self.channel_capacity); let (message_tx, agent_event_rx) = mpsc::channel::(self.channel_capacity); - let mut context = Context::new(messages, self.tool_definitions); + let mut context = Context::new(messages, Vec::new()); context.set_model_settings(self.model_settings); context.set_session_affinity_key(Some(self.session_affinity_key)); @@ -276,9 +272,9 @@ impl AgentBuilder { }; let agent = Agent::new(config, command_rx, message_tx); - let agent_handle = tokio::spawn(agent.run()); + let handle = AbortOnDropHandle::new(tokio::spawn(agent.run())); - Ok((command_tx, agent_event_rx, AgentHandle { handle: agent_handle })) + Ok((command_tx, agent_event_rx, AgentHandle { handle })) } } @@ -287,23 +283,27 @@ mod tests { use super::*; use crate::agent_spec::AgentSpecExposure; use crate::events::TurnEvent; + use crate::mcp::{McpConfig, ToolFilter}; use llm::ProviderConnectionOverrides; - use mcp_utils::client::ToolFilter; + use llm::testing::FakeLlmProvider; #[tokio::test] - async fn test_agent_handle_is_finished() { - let handle = AgentHandle { handle: tokio::spawn(async {}) }; + async fn closing_the_input_lets_the_agent_finish() { + let (tx, _rx, handle) = AgentBuilder::new(Arc::new(FakeLlmProvider::new(Vec::new()))).spawn().await.unwrap(); + + drop(tx); + handle.await_completion().await; } #[tokio::test] - async fn test_agent_handle_abort() { - let handle = AgentHandle { handle: tokio::spawn(std::future::pending::<()>()) }; - assert!(!handle.is_finished()); - handle.abort(); - while !handle.is_finished() { - tokio::task::yield_now().await; - } + async fn dropping_the_handle_stops_the_agent() { + let (_tx, mut rx, handle) = + AgentBuilder::new(Arc::new(FakeLlmProvider::new(Vec::new()))).spawn().await.unwrap(); + + drop(handle); + + assert!(rx.recv().await.is_none()); } #[tokio::test] @@ -348,13 +348,13 @@ mod tests { context_window: Some(200_000), prompts: vec![], provider_connections: ProviderConnectionOverrides::default(), - mcp_config_sources: Vec::new(), + mcp: McpConfig::default(), exposure: AgentSpecExposure::both(), tools: ToolFilter::default(), }; let dependencies = AgentDeps::default().with_session_affinity_key("conversation-123"); - let builder = AgentBuilder::from_spec(&spec, vec![], &dependencies).await.unwrap(); + let builder = AgentBuilder::from_spec(&spec, &dependencies).await.unwrap(); assert_eq!(builder.context_window, Some(200_000)); assert_eq!(builder.model_settings, settings); @@ -372,12 +372,12 @@ mod tests { context_window: None, prompts: vec![], provider_connections: ProviderConnectionOverrides::default(), - mcp_config_sources: Vec::new(), + mcp: McpConfig::default(), exposure: AgentSpecExposure::both(), tools: ToolFilter::default(), }; - let builder = AgentBuilder::from_spec(&spec, vec![], &AgentDeps::default()).await; + let builder = AgentBuilder::from_spec(&spec, &AgentDeps::default()).await; assert!(builder.is_ok()); } } diff --git a/crates/aether-core/src/core/agent_deps.rs b/crates/aether-core/src/core/agent_deps.rs index 35cecd5c0..f8e3b892c 100644 --- a/crates/aether-core/src/core/agent_deps.rs +++ b/crates/aether-core/src/core/agent_deps.rs @@ -1,7 +1,6 @@ use crate::core::AgentRegistry; use crate::events::{AgentObserver, DynObserverFactory, TraceContext}; use aether_auth::OAuthCredentialStorage; -use rmcp::model::ClientCapabilities; use std::sync::Arc; /// Cross-cutting dependencies threaded to every agent a run spawns — the root @@ -15,7 +14,6 @@ pub struct AgentDeps { /// that spawned them. pub parent_trace_context: Option, pub agent_registry: AgentRegistry, - pub mcp_client_capabilities: Option, pub session_affinity_key: Option, } @@ -38,23 +36,11 @@ impl AgentDeps { self } - pub fn with_mcp_client_capabilities(mut self, capabilities: ClientCapabilities) -> Self { - self.mcp_client_capabilities = Some(capabilities); - self - } - pub fn with_session_affinity_key(mut self, key: impl Into) -> Self { self.session_affinity_key = Some(key.into()); self } - pub fn supports_mcp_url_elicitation(&self) -> bool { - self.mcp_client_capabilities - .as_ref() - .and_then(|capabilities| capabilities.elicitation.as_ref()) - .is_some_and(|elicitation| elicitation.url.is_some()) - } - /// A fresh observer isolated to one agent, if a factory is configured. pub fn observer(&self, agent_name: &str) -> Option> { self.observer_factory @@ -62,22 +48,3 @@ impl AgentDeps { .map(|factory| factory.agent(Some(agent_name), self.parent_trace_context.as_ref())) } } - -#[cfg(test)] -mod tests { - use super::*; - use rmcp::model::{ElicitationCapability, FormElicitationCapability, UrlElicitationCapability}; - - #[test] - fn mcp_url_elicitation_support_requires_advertised_url_capability() { - assert!(!AgentDeps::default().supports_mcp_url_elicitation()); - - let mut form_only = ClientCapabilities::default(); - form_only.elicitation = Some(ElicitationCapability::new().with_form(FormElicitationCapability::new())); - assert!(!AgentDeps::default().with_mcp_client_capabilities(form_only).supports_mcp_url_elicitation()); - - let mut url = ClientCapabilities::default(); - url.elicitation = Some(ElicitationCapability::new().with_url(UrlElicitationCapability::new())); - assert!(AgentDeps::default().with_mcp_client_capabilities(url).supports_mcp_url_elicitation()); - } -} diff --git a/crates/aether-core/src/core/error.rs b/crates/aether-core/src/core/error.rs index 85e0f9bd0..a47d539a3 100644 --- a/crates/aether-core/src/core/error.rs +++ b/crates/aether-core/src/core/error.rs @@ -2,18 +2,12 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum AgentError { - /// MCP manager operation failed - #[error("MCP error: {0}")] - McpError(#[from] mcp_utils::client::McpError), /// LLM provider error #[error("LLM error: {0}")] LlmError(#[from] llm::LlmError), /// IO error (file operations, etc.) #[error("IO error: {0}")] IoError(String), - /// Generic error for other cases - #[error("{0}")] - Other(String), } pub type Result = std::result::Result; diff --git a/crates/aether-core/src/core/prompt.rs b/crates/aether-core/src/core/prompt.rs index 9bb9284ac..4e02194b8 100644 --- a/crates/aether-core/src/core/prompt.rs +++ b/crates/aether-core/src/core/prompt.rs @@ -1,4 +1,5 @@ use crate::core::{AgentError, Result}; +use crate::mcp::PROGRESSIVE_DISCOVERY_INSTRUCTIONS; use glob::glob; use schemars::{JsonSchema, Schema, SchemaGenerator}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; @@ -15,14 +16,14 @@ use utils::{PathOrObject, ResourcePath, is_false, string_or_object_schema}; #[derive(Debug, Clone, PartialEq)] pub enum Prompt { Text(String), - File { - path: PathBuf, - args: Option>, - cwd: PathBuf, - }, - /// MCP server instructions keyed by server name. `BTreeMap` gives a - /// stable render order, which is useful for prompt-caching - McpInstructions(BTreeMap), + File { path: PathBuf, args: Option>, cwd: PathBuf }, + McpInstructions(McpInstructions), +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct McpInstructions { + pub servers: BTreeMap, + pub progressive_discovery: bool, } /// Authored description of a prompt source — text, a file path, or a glob pattern. @@ -295,25 +296,22 @@ pub struct PromptCache { impl PromptCache { pub fn new(mut prompts: Vec) -> Self { if !prompts.iter().any(|p| matches!(p, Prompt::McpInstructions(_))) { - prompts.push(Prompt::McpInstructions(BTreeMap::new())); + prompts.push(Prompt::McpInstructions(McpInstructions::default())); } Self { prompts, entries: Vec::new() } } - pub fn update_mcp_instruction(&mut self, server: String, body: Option) { - for prompt in &mut self.prompts { - if let Prompt::McpInstructions(map) = prompt { - match body { - Some(text) => { - map.insert(server, text); - } - None => { - map.remove(&server); - } - } - return; - } + pub fn set_mcp_instructions(&mut self, instructions: McpInstructions) -> bool { + let Some(Prompt::McpInstructions(current)) = + self.prompts.iter_mut().find(|prompt| matches!(prompt, Prompt::McpInstructions(_))) + else { + return false; + }; + if *current == instructions { + return false; } + *current = instructions; + true } pub async fn render(&mut self) -> Result { @@ -339,20 +337,20 @@ impl PromptCache { } } -/// Format MCP instructions with XML tags for the system prompt. -fn format_mcp_instructions(instructions: &BTreeMap) -> String { - if instructions.is_empty() { - return String::new(); +fn format_mcp_instructions(instructions: &McpInstructions) -> String { + let mut sections = Vec::new(); + if !instructions.servers.is_empty() { + let mut parts = vec!["# MCP Server Instructions\n".to_string()]; + parts.push("You are connected to the following MCP servers:\n".to_string()); + for (server_name, body) in &instructions.servers { + parts.push(format!("\n{body}\n\n")); + } + sections.push(parts.join("\n")); } - - let mut parts = vec!["# MCP Server Instructions\n".to_string()]; - parts.push("You are connected to the following MCP servers:\n".to_string()); - - for (server_name, body) in instructions { - parts.push(format!("\n{body}\n\n")); + if instructions.progressive_discovery { + sections.push(format!("# Deferred MCP Tools\n\n{PROGRESSIVE_DISCOVERY_INSTRUCTIONS}")); } - - parts.join("\n") + sections.join("\n\n") } #[cfg(test)] @@ -419,14 +417,14 @@ mod tests { write(dir.path().join("AGENTS.md"), "cached body").unwrap(); let mut cache = PromptCache::new(vec![ Prompt::file(dir.path().join("AGENTS.md"), dir.path()), - Prompt::McpInstructions(BTreeMap::new()), + Prompt::McpInstructions(McpInstructions::default()), ]); cache.render().await.unwrap(); // Remove the source file to prove we cached things remove_file(dir.path().join("AGENTS.md")).unwrap(); - cache.update_mcp_instruction("srv".into(), Some("instr".into())); + cache.set_mcp_instructions(instructions(&[("srv", "instr")])); let rendered = cache.render().await.unwrap(); assert!(rendered.contains("cached body")); diff --git a/crates/aether-core/src/core/tool_execution.rs b/crates/aether-core/src/core/tool_execution.rs index 3a08a16b1..72f4c2b6d 100644 --- a/crates/aether-core/src/core/tool_execution.rs +++ b/crates/aether-core/src/core/tool_execution.rs @@ -1,5 +1,5 @@ -use crate::events::{SubAgentProgressPayload, TaskOutcome, TaskOutcomeState, ToolEvent, task_created_result}; -use crate::mcp::tool_bridge::{convert_tool_result, map_task_result_to_outcome}; +use crate::events::{SubAgentProgressPayload, TaskOutcome, ToolEvent, task_created_result}; +use crate::mcp::tool_bridge::{ToolRejection, convert_tool_result, map_task_result_to_outcome}; use llm::{ToolCallError, ToolCallRequest, ToolCallResult}; use mcp_utils::client::{CancellationToken, ToolCallEvent}; use rmcp::model::ProgressNotificationParam; @@ -17,8 +17,7 @@ pub(super) enum ToolExecutionUpdate { Event(ToolEvent), TaskCreated { result: ToolCallResult, event: ToolEvent }, Completed { result: Result, event: ToolEvent }, - TaskCompleted(TaskOutcome), - TaskCancelled(TaskOutcome), + TaskFinished(TaskOutcome), Retired, Ignored, } @@ -34,11 +33,10 @@ struct ToolExecution { phase: ToolExecutionPhase, } -#[derive(Clone, Copy, PartialEq, Eq)] +#[derive(PartialEq, Eq)] enum ToolExecutionPhase { Foreground, - Background, - Cancelling, + Background { task_id: String, cancelling: bool }, Retiring, } @@ -68,7 +66,10 @@ impl ToolExecutions { match event { ToolCallEvent::Progress(progress) => { let Some(execution) = self.executions.get(tool_id).filter(|execution| { - matches!(execution.phase, ToolExecutionPhase::Foreground | ToolExecutionPhase::Background) + matches!( + execution.phase, + ToolExecutionPhase::Foreground | ToolExecutionPhase::Background { cancelling: false, .. } + ) }) else { return ToolExecutionUpdate::Ignored; }; @@ -85,9 +86,9 @@ impl ToolExecutions { if execution.phase != ToolExecutionPhase::Foreground { return ToolExecutionUpdate::Ignored; } - execution.phase = ToolExecutionPhase::Background; - let request = execution.request.clone(); let task_id = task.task.task_id.clone(); + execution.phase = ToolExecutionPhase::Background { task_id: task_id.clone(), cancelling: false }; + let request = execution.request.clone(); ToolExecutionUpdate::TaskCreated { result: task_created_result(&request, &task_id), event: ToolEvent::TaskCreated { request, task_id, status_message: task.task.status_message }, @@ -97,7 +98,7 @@ impl ToolExecutions { let Some(execution) = self.executions.get(tool_id) else { return ToolExecutionUpdate::Ignored; }; - if execution.phase != ToolExecutionPhase::Background { + if !matches!(execution.phase, ToolExecutionPhase::Background { cancelling: false, .. }) { return ToolExecutionUpdate::Ignored; } ToolExecutionUpdate::Event(ToolEvent::TaskStatus { @@ -107,53 +108,40 @@ impl ToolExecutions { status_message: task.status_message, }) } - ToolCallEvent::TaskComplete { task, result } => { - if self.take_retiring(tool_id).is_some() { - return ToolExecutionUpdate::Retired; - } - let Some(execution) = self.take_background(tool_id) else { - return ToolExecutionUpdate::Ignored; - }; - ToolExecutionUpdate::TaskCompleted(map_task_result_to_outcome( - execution.request, - task, - result, - &self.spill_dir, - )) - } - ToolCallEvent::Cancelled { task_id } => { - if self.take_retiring(tool_id).is_some() { - return ToolExecutionUpdate::Retired; - } - let Some(execution) = self.take_background(tool_id) else { + ToolCallEvent::Done(result) => { + let Some(ToolExecution { request, phase, .. }) = self.executions.remove(tool_id) else { return ToolExecutionUpdate::Ignored; }; - ToolExecutionUpdate::TaskCancelled(TaskOutcome { - request: execution.request, - task_id: task_id.unwrap_or_else(|| UNASSIGNED_TASK_ID.to_string()), - state: TaskOutcomeState::Cancelled, - }) - } - ToolCallEvent::Complete(outcome) => { - if self.take_retiring(tool_id).is_some() { - return ToolExecutionUpdate::Retired; - } - let Some(execution) = self.take_foreground(tool_id) else { - return ToolExecutionUpdate::Ignored; - }; - match convert_tool_result(&execution.request, outcome, &self.spill_dir) { - Ok((result, result_meta)) => ToolExecutionUpdate::Completed { - result: Ok(result.clone()), - event: ToolEvent::Result { result, result_meta }, + match phase { + ToolExecutionPhase::Retiring => ToolExecutionUpdate::Retired, + ToolExecutionPhase::Foreground => match convert_tool_result(&request, result, &self.spill_dir) { + Ok((result, result_meta)) => ToolExecutionUpdate::Completed { + result: Ok(result.clone()), + event: ToolEvent::Result { result, result_meta }, + }, + Err(error) => ToolExecutionUpdate::Completed { + result: Err(error.clone()), + event: ToolEvent::Error { error }, + }, }, - Err(error) => { - ToolExecutionUpdate::Completed { result: Err(error.clone()), event: ToolEvent::Error { error } } + ToolExecutionPhase::Background { task_id, .. } => { + let outcome = map_task_result_to_outcome(request, task_id, result, &self.spill_dir); + ToolExecutionUpdate::TaskFinished(outcome) } } } } } + /// Fails a tool call that never reached a server. + pub(super) fn reject(&mut self, tool_id: &str, rejection: &ToolRejection) -> ToolExecutionUpdate { + let Some(execution) = self.executions.remove(tool_id) else { + return ToolExecutionUpdate::Ignored; + }; + let error = ToolCallError::from_request(&execution.request, rejection.to_string()); + ToolExecutionUpdate::Completed { result: Err(error.clone()), event: ToolEvent::Error { error } } + } + pub(super) fn retire_foreground(&mut self) { for execution in self.executions.values_mut() { if execution.phase == ToolExecutionPhase::Foreground { @@ -164,13 +152,13 @@ impl ToolExecutions { pub(super) fn abort(&mut self, policy: &ToolAbortPolicy) -> Vec { let mut removed = Vec::new(); - self.executions.retain(|tool_id, execution| match execution.phase { - ToolExecutionPhase::Background | ToolExecutionPhase::Cancelling => { + self.executions.retain(|tool_id, execution| match &mut execution.phase { + ToolExecutionPhase::Background { cancelling, .. } => { execution.cancellation_token.cancel(); - execution.phase = match policy { - ToolAbortPolicy::PreserveBackgroundAcknowledgements => ToolExecutionPhase::Cancelling, - ToolAbortPolicy::CancelAll => ToolExecutionPhase::Retiring, - }; + match policy { + ToolAbortPolicy::PreserveBackgroundAcknowledgements => *cancelling = true, + ToolAbortPolicy::CancelAll => execution.phase = ToolExecutionPhase::Retiring, + } true } ToolExecutionPhase::Retiring if matches!(policy, ToolAbortPolicy::CancelAll) => true, @@ -182,32 +170,8 @@ impl ToolExecutions { }); removed } - - fn take_retiring(&mut self, tool_id: &str) -> Option { - if self.executions.get(tool_id)?.phase != ToolExecutionPhase::Retiring { - return None; - } - self.executions.remove(tool_id) - } - - fn take_foreground(&mut self, tool_id: &str) -> Option { - if self.executions.get(tool_id)?.phase != ToolExecutionPhase::Foreground { - return None; - } - self.executions.remove(tool_id) - } - - fn take_background(&mut self, tool_id: &str) -> Option { - let phase = self.executions.get(tool_id)?.phase; - if !matches!(phase, ToolExecutionPhase::Background | ToolExecutionPhase::Cancelling) { - return None; - } - self.executions.remove(tool_id) - } } -const UNASSIGNED_TASK_ID: &str = "pending"; - /// Tools smuggle structured updates through the progress message string. Decode /// them once here so every downstream consumer sees typed events. fn progress_event(request: ToolCallRequest, progress: ProgressNotificationParam) -> ToolEvent { diff --git a/crates/mcp-utils/src/docs/mcp_server_config.md b/crates/aether-core/src/docs/mcp_server_config.md similarity index 100% rename from crates/mcp-utils/src/docs/mcp_server_config.md rename to crates/aether-core/src/docs/mcp_server_config.md diff --git a/crates/aether-core/src/docs/tool_filter.md b/crates/aether-core/src/docs/tool_filter.md index 3f8f12710..9f288b860 100644 --- a/crates/aether-core/src/docs/tool_filter.md +++ b/crates/aether-core/src/docs/tool_filter.md @@ -25,6 +25,6 @@ Deny always wins over allow: Annotation matcher fields are `readOnly`, `destructive`, `idempotent`, and `openWorld`. Missing tool annotations or missing annotation fields do not receive default boolean values and do not match. MCP name matchers are evaluated against the canonical namespaced name -(`server__tool`), not the server-local name. The decision is stored in the -per-agent `ToolCatalog` and applies consistently to model-visible definitions, +(`server__tool`), not the server-local name. The agent's `McpCatalog` applies the +filter on top of the MCP gateway, consistently across model-visible definitions, deferred discovery, instructions, and both model-visible and deferred execution routes. diff --git a/crates/aether-core/src/events/user_message.rs b/crates/aether-core/src/events/user_message.rs index 25d1c65b4..aa37431c3 100644 --- a/crates/aether-core/src/events/user_message.rs +++ b/crates/aether-core/src/events/user_message.rs @@ -1,4 +1,4 @@ -use llm::{ChatMessage, ContentBlock, MessageId, ReasoningEffort, StreamingModelProvider, ToolDefinition}; +use llm::{ChatMessage, ContentBlock, MessageId, ReasoningEffort, StreamingModelProvider}; /// The unified command type sent to the agent input channel. /// @@ -56,14 +56,8 @@ impl std::fmt::Debug for UserCommand { } /// Runtime-internal operations that modify agent state without user interaction. -/// -/// These are sent by the runtime controller or MCP event pumps to update -/// tools, model, instructions, or to sync conversation history across -/// agent runtimes. pub enum AgentCommand { SwitchModel(Box), - UpdateTools(Vec), - UpdateMcpInstructions { server: String, body: Option }, SetReasoningEffort(Option), ReplaceConversation(Vec), } @@ -74,12 +68,6 @@ impl std::fmt::Debug for AgentCommand { AgentCommand::SwitchModel(provider) => { f.debug_tuple("SwitchModel").field(&provider.display_name()).finish() } - AgentCommand::UpdateTools(tools) => f.debug_tuple("UpdateTools").field(&tools.len()).finish(), - AgentCommand::UpdateMcpInstructions { server, body } => f - .debug_struct("UpdateMcpInstructions") - .field("server", server) - .field("body_len", &body.as_ref().map(String::len)) - .finish(), AgentCommand::SetReasoningEffort(effort) => f.debug_tuple("SetReasoningEffort").field(effort).finish(), AgentCommand::ReplaceConversation(messages) => { f.debug_tuple("ReplaceConversation").field(&messages.len()).finish() diff --git a/crates/aether-core/src/mcp/config.rs b/crates/aether-core/src/mcp/config.rs new file mode 100644 index 000000000..411e3036d --- /dev/null +++ b/crates/aether-core/src/mcp/config.rs @@ -0,0 +1,264 @@ +use super::ToolExposure; +use aether_auth::OAuthClientRegistration; +use http::header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue, InvalidHeaderName, InvalidHeaderValue}; +use mcp_utils::client::{OAuthSettings, Transport}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, HashMap}; +use std::num::NonZeroU16; +use std::path::Path; +use utils::variables::{VarError, Vars}; + +#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)] +pub struct McpConfig { + #[serde(alias = "mcpServers")] + pub servers: BTreeMap, +} + +#[doc = include_str!("../docs/mcp_server_config.md")] +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(untagged)] +pub enum McpServerConfig { + Stdio(StdioServerConfig), + Remote(RemoteServerConfig), + InMemory(InMemoryServerConfig), +} + +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct StdioServerConfig { + #[serde(rename = "type", default)] + pub type_: StdioType, + + pub command: String, + + #[serde(default)] + pub args: Vec, + + #[serde(default)] + pub env: HashMap, + + #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")] + pub defer_tools: ToolExposure, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct McpOAuthConfig { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_metadata_url: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub callback_port: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct RemoteServerConfig { + #[serde(rename = "type")] + pub type_: RemoteType, + + pub url: String, + + #[serde(default)] + pub headers: HashMap, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub oauth: Option, + + #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")] + pub defer_tools: ToolExposure, +} + +#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct InMemoryServerConfig { + #[serde(rename = "type")] + pub type_: InMemoryType, + + #[serde(default)] + pub args: Vec, + + #[serde(rename = "deferTools", default, skip_serializing_if = "is_model_visible")] + pub defer_tools: ToolExposure, +} + +#[derive(Debug)] +pub(super) enum ServerSource { + Ready(Transport), + InMemory(Vec), +} + +#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, JsonSchema, PartialEq)] +pub enum StdioType { + #[default] + #[serde(rename = "stdio")] + Stdio, +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)] +pub enum RemoteType { + #[serde(rename = "http")] + Http, +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)] +pub enum InMemoryType { + #[serde(rename = "in-memory")] + InMemory, +} + +#[derive(Debug, thiserror::Error)] +pub enum ParseError { + #[error("Failed to read config file: {0}")] + IoError(#[from] std::io::Error), + + #[error("Invalid JSON: {0}")] + JsonError(#[from] serde_json::Error), + + #[error("Variable expansion failed: {0}")] + VarError(#[from] VarError), + + #[error("Invalid HTTP header name {name:?}: {source}")] + InvalidHeaderName { + name: String, + #[source] + source: InvalidHeaderName, + }, + + #[error("Invalid HTTP header value for {name:?}: {source}")] + InvalidHeaderValue { + name: String, + #[source] + source: InvalidHeaderValue, + }, +} + +pub fn http_transport(url: String, headers: HeaderMap, oauth: Option) -> Transport { + let oauth = (!headers.contains_key(AUTHORIZATION)).then(|| oauth.unwrap_or_default().into_settings()); + Transport::Http { url, headers, oauth } +} + +impl McpConfig { + pub fn new(servers: BTreeMap) -> Self { + Self { servers } + } + + pub fn from_json_file(path: impl AsRef) -> Result { + let content = std::fs::read_to_string(path)?; + Self::from_json(&content) + } + + pub fn from_json(json: &str) -> Result { + Ok(serde_json::from_str(json)?) + } + + pub fn is_empty(&self) -> bool { + self.servers.is_empty() + } + + pub fn defer_all_tools(&mut self) { + for exposure in self.servers.values_mut().map(McpServerConfig::defer_tools_mut) { + if !exposure.defers_tools() { + *exposure = ToolExposure::deferred_all(); + } + } + } +} + +/// Merges configs in order; a server named again replaces the earlier definition. +impl FromIterator for McpConfig { + fn from_iter>(configs: T) -> Self { + Self::new(configs.into_iter().flat_map(|config| config.servers).collect()) + } +} + +impl McpServerConfig { + pub fn defer_tools(&self) -> &ToolExposure { + match self { + McpServerConfig::Stdio(config) => &config.defer_tools, + McpServerConfig::Remote(config) => &config.defer_tools, + McpServerConfig::InMemory(config) => &config.defer_tools, + } + } + + pub(super) fn resolve(self, vars: &Vars, cwd: &Path) -> Result { + Ok(match self { + McpServerConfig::Stdio(config) => ServerSource::Ready(config.into_transport(vars, cwd)?), + McpServerConfig::Remote(config) => ServerSource::Ready(config.into_transport(vars)?), + McpServerConfig::InMemory(config) => ServerSource::InMemory(expand_all(vars, &config.args)?), + }) + } + + fn defer_tools_mut(&mut self) -> &mut ToolExposure { + match self { + McpServerConfig::Stdio(config) => &mut config.defer_tools, + McpServerConfig::Remote(config) => &mut config.defer_tools, + McpServerConfig::InMemory(config) => &mut config.defer_tools, + } + } +} + +impl StdioServerConfig { + pub fn into_transport(self, vars: &Vars, cwd: &Path) -> Result { + let env = self.env.into_iter().map(|(name, value)| Ok((name, vars.expand(&value)?))); + Ok(Transport::Stdio { + command: vars.expand(&self.command)?, + args: expand_all(vars, &self.args)?, + env: env.collect::>()?, + cwd: Some(cwd.to_path_buf()), + }) + } +} + +impl McpOAuthConfig { + fn into_settings(self) -> OAuthSettings { + let callback_port = self.callback_port.map_or(AETHER_OAUTH_CALLBACK_PORT, NonZeroU16::get); + let registration = match self { + Self { client_id: Some(client_id), .. } => OAuthClientRegistration::PreRegistered(client_id), + Self { client_metadata_url: Some(url), .. } => OAuthClientRegistration::ClientMetadata(url), + Self { .. } => OAuthClientRegistration::ClientMetadata(AETHER_OAUTH_CLIENT_METADATA_URL.to_string()), + }; + OAuthSettings { registration, callback_port } + } +} + +impl RemoteServerConfig { + pub fn into_transport(self, vars: &Vars) -> Result { + let mut headers = HeaderMap::with_capacity(self.headers.len()); + for (name, value) in self.headers { + let header_name = name + .parse::() + .map_err(|source| ParseError::InvalidHeaderName { name: name.clone(), source })?; + let mut header_value = vars + .expand(&value)? + .parse::() + .map_err(|source| ParseError::InvalidHeaderValue { name, source })?; + header_value.set_sensitive(true); + headers.insert(header_name, header_value); + } + let oauth = self.oauth.map(|oauth| expand_oauth(vars, oauth)).transpose()?; + Ok(http_transport(vars.expand(&self.url)?, headers, oauth)) + } +} + +const AETHER_OAUTH_CLIENT_METADATA_URL: &str = "https://aether-agent.io/oauth/client-metadata.json"; +const AETHER_OAUTH_CALLBACK_PORT: u16 = 3118; + +fn is_model_visible(exposure: &ToolExposure) -> bool { + !exposure.defers_tools() +} + +fn expand_all(vars: &Vars, values: &[String]) -> Result, VarError> { + values.iter().map(|value| vars.expand(value)).collect() +} + +fn expand_oauth(vars: &Vars, oauth: McpOAuthConfig) -> Result { + let expand = |value: Option| value.map(|value| vars.expand(&value)).transpose(); + Ok(McpOAuthConfig { + client_id: expand(oauth.client_id)?, + client_metadata_url: expand(oauth.client_metadata_url)?, + callback_port: oauth.callback_port, + }) +} diff --git a/crates/aether-core/src/mcp/deferred_tools.rs b/crates/aether-core/src/mcp/deferred_tools.rs new file mode 100644 index 000000000..ffbe770f8 --- /dev/null +++ b/crates/aether-core/src/mcp/deferred_tools.rs @@ -0,0 +1,194 @@ +use super::mcp_catalog::{CatalogCore, Tier, namespaced}; +use mcp_utils::client::{ClientError, ClientOptions, McpClient, ToolCall, ToolCallError, ToolCallOptions, Transport}; +use mcp_utils::server::{McpServer, ServerHandle, UnixSocket}; +use rmcp::model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ErrorData, Implementation, ListToolsResult, + PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, +}; +use rmcp::{RoleServer, ServerHandler, service::RequestContext}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use std::future::{Future, ready}; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use thiserror::Error; + +pub const AETHER_MCP_IPC_SOCKET: &str = "AETHER_MCP_IPC_SOCKET"; +pub(crate) const PROGRESSIVE_DISCOVERY_INSTRUCTIONS: &str = include_str!("progressive_discovery_instructions.md"); + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeferredServer { + pub name: String, + pub description: String, + pub tools: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeferredTool { + pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, +} + +pub struct DeferredToolsClient { + client: McpClient, +} + +#[derive(Debug, Error)] +pub enum DeferredToolsError { + #[error("failed to connect to the active Aether session: {0}")] + Connect(#[source] ClientError), + #[error("deferred tool discovery timed out after {} seconds", .0.as_secs())] + TimedOut(Duration), + #[error("failed to list deferred servers: {0}")] + ListServers(#[source] ToolCallError), + #[error("invalid deferred server listing: {0}")] + MalformedListing(#[source] serde_json::Error), + #[error("failed to list deferred tools: {0}")] + ListTools(#[source] ClientError), +} + +impl DeferredToolsClient { + pub async fn connect(socket: PathBuf) -> Result { + let options = ClientOptions::default().implementation(Implementation::new("aether", env!("CARGO_PKG_VERSION"))); + let client = McpClient::connect("aether", Transport::Unix(socket), &options).await; + client.map(|client| Self { client }).map_err(DeferredToolsError::Connect) + } + + pub async fn servers(&self) -> Result, DeferredToolsError> { + let options = ToolCallOptions::with_timeout(DISCOVERY_TIMEOUT); + let listed = self.client.call_tool(LIST_SERVERS_TOOL, Map::new(), options).result().await; + let listed = listed.map_err(|error| match error { + ToolCallError::TimedOut(timeout) => DeferredToolsError::TimedOut(timeout), + error => DeferredToolsError::ListServers(error), + })?; + serde_json::from_value(listed.structured_content.unwrap_or_default()) + .map_err(DeferredToolsError::MalformedListing) + } + + pub async fn tool(&self, server: &str, tool: &str) -> Result, DeferredToolsError> { + let tools = tokio::time::timeout(DISCOVERY_TIMEOUT, self.client.list_tools()) + .await + .map_err(|_| DeferredToolsError::TimedOut(DISCOVERY_TIMEOUT))? + .map_err(DeferredToolsError::ListTools)?; + let name = namespaced(server, tool); + Ok(tools.into_iter().find(|definition| definition.name == name)) + } + + pub fn call(&self, server: &str, tool: &str, args: Map, timeout: Duration) -> ToolCall { + self.client.call_tool(&namespaced(server, tool), args, ToolCallOptions::with_timeout(timeout)) + } + + pub async fn close(&self) { + self.client.close().await; + } +} + +pub(super) fn serve(catalog: CatalogCore, socket: UnixSocket) -> ServerHandle { + McpServer::new(DeferredToolsMcp { catalog }).serve_unix(socket) +} + +const LIST_SERVERS_TOOL: &str = "_aether_list_servers"; +const CALL_TIMEOUT: Duration = Duration::from_mins(10); +const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); + +struct DeferredToolsMcp { + catalog: CatalogCore, +} + +impl DeferredToolsMcp { + fn tools(&self) -> Vec { + let list_servers = + Tool::new(LIST_SERVERS_TOOL, "List connected MCP servers with deferred tools", Arc::new(Map::new())); + let mut tools: Vec = self.catalog.tools(Tier::Deferred).into_iter().map(|(_, tool)| tool).collect(); + tools.push(list_servers); + tools + } + + fn list_servers(&self) -> CallToolResult { + let servers = self + .catalog + .servers(Tier::Deferred) + .into_iter() + .map(|(server, tools)| DeferredServer { + name: server.name().to_string(), + description: server.description().unwrap_or_else(|| server.name().to_string()), + tools: tools + .into_iter() + .map(|tool| DeferredTool { + name: tool.name.to_string(), + description: tool.description.map(Into::into), + }) + .collect(), + }) + .collect::>(); + CallToolResult::structured(serde_json::to_value(servers).expect("deferred servers serialize to JSON")) + } +} + +impl ServerHandler for DeferredToolsMcp { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("aether-deferred-tools", env!("CARGO_PKG_VERSION"))) + } + + fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + ready(Ok(ListToolsResult::with_all_items(self.tools()))) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + if request.name == LIST_SERVERS_TOOL { + return Ok(self.list_servers().into()); + } + let name = &request.name; + if self.catalog.tier_of(name) == Some(Tier::Direct) { + let message = format!("Tool '{name}' is exposed directly; call '{name}' as a regular tool instead"); + return Err(ErrorData::invalid_params(message, None)); + } + + let options = ToolCallOptions { timeout: Some(CALL_TIMEOUT), meta: request.meta, cancel: context.ct.clone() }; + let args = request.arguments.unwrap_or_default(); + let Some(call) = self.catalog.call(Tier::Deferred, name, args, options) else { + return Err(ErrorData::invalid_params(format!("Tool not found: {name}"), None)); + }; + forward(call, name, &context).await + } +} + +async fn forward( + call: ToolCall, + name: &str, + context: &RequestContext, +) -> Result { + let started = Instant::now(); + let progress_token = context.meta.get_progress_token(); + let result = call + .result_with_progress(|mut progress| { + let (peer, token) = (context.peer.clone(), progress_token.clone()); + async move { + if let Some(token) = token { + progress.progress_token = token; + let _ = peer.notify_progress(progress).await; + } + } + }) + .await; + + let outcome = match &result { + Ok(_) => "success", + Err(ToolCallError::Cancelled) => "cancelled", + Err(_) => "error", + }; + tracing::info!(tool = name, outcome, duration_ms = started.elapsed().as_millis(), "MCP tool call forwarded"); + + result.map(Into::into).map_err(|error| ErrorData::internal_error(error.to_string(), None)) +} diff --git a/crates/aether-core/src/mcp/exposure.rs b/crates/aether-core/src/mcp/exposure.rs new file mode 100644 index 000000000..2c99a46c6 --- /dev/null +++ b/crates/aether-core/src/mcp/exposure.rs @@ -0,0 +1,61 @@ +use super::{ToolFilter, ToolMatcher}; +use rmcp::model::Tool; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize, JsonSchema)] +#[serde(from = "ToolExposureConfig", into = "ToolExposureConfig")] +#[schemars(with = "ToolExposureConfig")] +pub enum ToolExposure { + #[default] + ModelVisible, + Deferred(ToolFilter), +} + +impl ToolExposure { + pub fn deferred_all() -> Self { + Self::Deferred(ToolFilter::default()) + } + + pub fn defers_tools(&self) -> bool { + matches!(self, Self::Deferred(_)) + } + + pub fn is_deferred(&self, tool: &Tool) -> bool { + matches!(self, Self::Deferred(filter) if filter.allows(&tool.name, tool)) + } +} + +#[derive(Deserialize, Serialize, JsonSchema)] +#[serde(untagged, deny_unknown_fields)] +enum ToolExposureConfig { + Enabled(bool), + Rules { + #[serde(default, skip_serializing_if = "Vec::is_empty")] + include: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + exclude: Vec, + }, +} + +impl From for ToolExposure { + fn from(repr: ToolExposureConfig) -> Self { + match repr { + ToolExposureConfig::Enabled(false) => Self::ModelVisible, + ToolExposureConfig::Enabled(true) => Self::deferred_all(), + ToolExposureConfig::Rules { include, exclude } => { + Self::Deferred(ToolFilter { allow: include, deny: exclude }) + } + } + } +} + +impl From for ToolExposureConfig { + fn from(exposure: ToolExposure) -> Self { + match exposure { + ToolExposure::ModelVisible => Self::Enabled(false), + ToolExposure::Deferred(filter) if filter.is_empty() => Self::Enabled(true), + ToolExposure::Deferred(ToolFilter { allow, deny }) => Self::Rules { include: allow, exclude: deny }, + } + } +} diff --git a/crates/aether-core/src/mcp/gateway_service.rs b/crates/aether-core/src/mcp/gateway_service.rs deleted file mode 100644 index 08d081cd9..000000000 --- a/crates/aether-core/src/mcp/gateway_service.rs +++ /dev/null @@ -1,187 +0,0 @@ -use super::McpHandle; -use futures::StreamExt; -use mcp_utils::client::{CallToolOptions, CancellationToken, ToolCallEvent, ToolRoute}; -use mcp_utils::tool_gateway::LIST_SERVERS_TOOL; -use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ErrorData, Implementation, ListToolsResult, - PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool, ToolAnnotations, -}; -use rmcp::{RoleServer, ServerHandler, service::RequestContext}; -use serde_json::{Map, json}; -use std::{ - sync::Arc, - time::{Duration, Instant}, -}; - -#[derive(Clone)] -pub struct GatewayService { - handle: McpHandle, -} - -impl GatewayService { - pub fn new(handle: McpHandle) -> Self { - Self { handle } - } - - fn tools(&self) -> Vec { - let snapshot = self.handle.snapshot(); - let mut tools = snapshot - .catalog() - .tools() - .deferred - .into_iter() - .map(|tool| { - let definition = tool.definition(); - let schema = definition.parameters.as_object().cloned().unwrap_or_default(); - let mut result = Tool::new(definition.name.clone(), definition.description.clone(), Arc::new(schema)); - if let Some(annotations) = &definition.annotations { - let mut converted = ToolAnnotations::new(); - converted.title.clone_from(&annotations.title); - converted.read_only_hint = annotations.read_only_hint; - converted.destructive_hint = annotations.destructive_hint; - converted.idempotent_hint = annotations.idempotent_hint; - converted.open_world_hint = annotations.open_world_hint; - result = result.with_annotations(converted); - } - result - }) - .collect::>(); - tools.push(Tool::new( - LIST_SERVERS_TOOL, - "List connected MCP servers with deferred tools", - Arc::new(Map::new()), - )); - tools - } - - fn list_servers(&self) -> CallToolResult { - let servers = self.handle.snapshot().catalog().discoverable_deferred_servers(); - let value = serde_json::to_value( - servers - .iter() - .map(|server| json!({ "name": server.name, "description": server.description })) - .collect::>(), - ) - .expect("server descriptions serialize"); - CallToolResult::structured(value) - } -} - -impl ServerHandler for GatewayService { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("aether-deferred-tool-gateway", env!("CARGO_PKG_VERSION"))) - } - - fn list_tools( - &self, - _request: Option, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(Ok(ListToolsResult::with_all_items(self.tools()))) - } - - async fn call_tool( - &self, - request: CallToolRequestParams, - context: RequestContext, - ) -> Result { - if request.name == LIST_SERVERS_TOOL { - return Ok(self.list_servers().into()); - } - - let (server, tool) = request - .name - .split_once("__") - .ok_or_else(|| ErrorData::invalid_params("deferred tool names must use server__tool", None))?; - - let arguments = request.arguments.unwrap_or_default(); - let cancellation = CancellationToken::new(); - let mut guard = CallCancellationGuard::new(cancellation.clone()); - let started = Instant::now(); - let mut events = self.handle.call( - ToolRoute::Deferred { server: server.to_string(), tool: tool.to_string() }, - arguments, - CallToolOptions { timeout: Duration::from_mins(10), meta: request.meta, cancel: cancellation.clone() }, - ); - let peer = context.peer; - let disconnected = async move { - while !peer.is_transport_closed() { - tokio::time::sleep(Duration::from_millis(25)).await; - } - }; - tokio::pin!(disconnected); - - loop { - let event = tokio::select! { - event = events.next() => event, - () = context.ct.cancelled() => None, - () = &mut disconnected => None, - }; - let Some(event) = event else { - cancellation.cancel(); - tracing::info!( - route = "deferred", - server, - tool, - outcome = "cancelled", - duration_ms = started.elapsed().as_millis(), - "deferred MCP gateway call completed" - ); - guard.disarm(); - return Err(ErrorData::internal_error("deferred tool call was cancelled", None)); - }; - match event { - ToolCallEvent::Complete(result) | ToolCallEvent::TaskComplete { result, .. } => { - let outcome = if result.is_ok() { "success" } else { "error" }; - tracing::info!( - route = "deferred", - server, - tool, - outcome, - duration_ms = started.elapsed().as_millis(), - "deferred MCP gateway call completed" - ); - guard.disarm(); - return result.map(Into::into).map_err(|error| ErrorData::internal_error(error.to_string(), None)); - } - ToolCallEvent::Cancelled { .. } => { - tracing::info!( - route = "deferred", - server, - tool, - outcome = "cancelled", - duration_ms = started.elapsed().as_millis(), - "deferred MCP gateway call completed" - ); - guard.disarm(); - return Err(ErrorData::internal_error("deferred tool call was cancelled", None)); - } - ToolCallEvent::Progress(_) | ToolCallEvent::TaskCreated(_) | ToolCallEvent::TaskStatus(_) => {} - } - } - } -} - -struct CallCancellationGuard { - token: CancellationToken, - armed: bool, -} - -impl CallCancellationGuard { - fn new(token: CancellationToken) -> Self { - Self { token, armed: true } - } - - fn disarm(&mut self) { - self.armed = false; - } -} - -impl Drop for CallCancellationGuard { - fn drop(&mut self) { - if self.armed { - self.token.cancel(); - } - } -} diff --git a/crates/aether-core/src/mcp/mcp_builder.rs b/crates/aether-core/src/mcp/mcp_builder.rs index 8dcbd087b..cb2cd07c7 100644 --- a/crates/aether-core/src/mcp/mcp_builder.rs +++ b/crates/aether-core/src/mcp/mcp_builder.rs @@ -1,32 +1,16 @@ -use mcp_utils::client::{ - InMemoryServerSpec, McpClientEvent, McpConfig, McpConnectionDetails, McpError, McpManager, McpServer, McpTransport, - OAuthHandlerFactory, PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME, ParseError, RuntimeMcpServer, RuntimeMcpTransport, - ToolFilter, -}; -use mcp_utils::tool_gateway::{AETHER_MCP_IPC_SOCKET, UnixSocketMcpTransport, UnixSocketPath, UnixSocketServer}; +use mcp_utils::client::{ClientOptions, ElicitationSink, Transport}; +use mcp_utils::gateway::{McpGateway, ServerSpec}; +use mcp_utils::model::Implementation; +use mcp_utils::server::{McpServer, ServeError, UnixSocket}; use utils::{SettingsStore, variables::Vars}; -use crate::agent_spec::McpConfigSource; +use super::config::ServerSource; +use super::{McpCatalog, McpConfig, McpServerConfig, ParseError, ToolExposure, ToolFilter}; use crate::core::AgentDeps; -use crate::events::{AgentCommand, Command}; - -use super::{ - gateway_service::GatewayService, - mcp_handle::McpHandle, - run_mcp_task::{ManagerCommand, run_mcp_task}, -}; -use futures::future::BoxFuture; -use rmcp::{RoleServer, service::DynService}; -use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::Arc; -use tokio::{ - sync::{ - mpsc::{self, Receiver}, - watch, - }, - task::JoinHandle, -}; +use thiserror::Error; pub fn mcp(root_dir: impl AsRef) -> McpBuilder { McpBuilder::new(root_dir) @@ -34,159 +18,31 @@ pub fn mcp(root_dir: impl AsRef) -> McpBuilder { #[derive(Clone)] pub struct RuntimeServices { - pub mcp: McpHandle, pub root_dir: PathBuf, pub agent_deps: AgentDeps, - pub shell_environment: BTreeMap, -} - -pub type ServerFactory = Box< - dyn Fn(InMemoryServerSpec, RuntimeServices) -> BoxFuture<'static, Box>> + Send + Sync, ->; - -/// Owns a spawned MCP manager. Dropping this value aborts the manager task. -pub struct McpRuntime { - mcp: McpHandle, - handle: JoinHandle<()>, - agent_sync_handle: Option>, - gateway: Option, -} - -impl McpRuntime { - pub async fn shutdown(&mut self) { - self.handle.abort(); - if let Some(handle) = self.agent_sync_handle.take() { - handle.abort(); - let _ = handle.await; - } - let _ = (&mut self.handle).await; - self.gateway.take(); - } - - pub fn handle(&self) -> &McpHandle { - &self.mcp - } - - pub fn gateway_endpoint(&self) -> Option<&Path> { - self.gateway.as_ref().map(UnixSocketServer::path) - } -} - -impl Drop for McpRuntime { - fn drop(&mut self) { - self.handle.abort(); - if let Some(handle) = &self.agent_sync_handle { - handle.abort(); - } - } -} - -/// A freshly spawned MCP manager paired with its event stream. Consumers -/// receive incremental updates over the event stream (starting with an initial -/// `ServerStatusesChanged` reflecting every configured server in `Connecting`) -/// and can call [`split`](Self::split) to separate the stream from the -/// [`McpRuntime`] that keeps the manager task alive. -pub struct McpSession { - runtime: McpRuntime, - event_rx: Receiver, -} - -impl McpSession { - pub fn handle(&self) -> &McpHandle { - self.runtime.handle() - } - - pub fn gateway_endpoint(&self) -> Option<&Path> { - self.runtime.gateway_endpoint() - } - - /// Synchronize this session's current and future tools and instructions with - /// one agent. Initial state is sent before this method returns. - pub async fn connect_agent(mut self, agent_tx: mpsc::Sender) -> Self { - assert!(self.runtime.agent_sync_handle.is_none(), "an MCP session can only connect one agent"); - let mut snapshots = self.runtime.handle().subscribe(); - let initial = snapshots.borrow_and_update().clone(); - let mut previous_tools = initial.tool_definitions(); - let mut previous_instructions = initial.model_instructions(); - if agent_tx.send(Command::agent(AgentCommand::UpdateTools(previous_tools.clone()))).await.is_err() { - return self; - } - for (server, body) in &previous_instructions { - if agent_tx - .send(Command::agent(AgentCommand::UpdateMcpInstructions { - server: server.clone(), - body: Some(body.clone()), - })) - .await - .is_err() - { - return self; - } - } - - let agent_tx = agent_tx.downgrade(); - self.runtime.agent_sync_handle = Some(tokio::spawn(async move { - while snapshots.changed().await.is_ok() { - let Some(agent_tx) = agent_tx.upgrade() else { - break; - }; - let snapshot = snapshots.borrow_and_update().clone(); - let tools = snapshot.tool_definitions(); - if tools != previous_tools { - if agent_tx.send(Command::agent(AgentCommand::UpdateTools(tools.clone()))).await.is_err() { - break; - } - previous_tools = tools; - } - - let instructions = snapshot.model_instructions(); - let servers = previous_instructions.keys().chain(instructions.keys()).cloned().collect::>(); - for server in servers { - let previous = previous_instructions.get(&server); - let next = instructions.get(&server); - if previous != next - && agent_tx - .send(Command::agent(AgentCommand::UpdateMcpInstructions { server, body: next.cloned() })) - .await - .is_err() - { - return; - } - } - previous_instructions = instructions; - } - })); - self - } - - /// Block until the manager finishes bootstrapping every initially-configured - /// server, then return the consolidated snapshot. Returns `None` if the - /// event channel closes before `ConnectionReady` is received. - pub async fn block_until_ready(&mut self) -> Option { - while let Some(event) = self.event_rx.recv().await { - if let McpClientEvent::ConnectionReady(snapshot) = event { - return Some(snapshot); - } - } - None - } - - pub fn split(self) -> (McpRuntime, Receiver) { - (self.runtime, self.event_rx) - } + pub deferred_tools_socket: Option, } pub struct McpBuilder { - servers: Vec, + servers: Vec, factories: HashMap, - mcp_channel_capacity: usize, root_dir: PathBuf, - oauth_handler_factory: Option, agent_deps: AgentDeps, - aether_home: Option, vars: Vars, tool_filter: ToolFilter, - progressive_discovery_instructions: Option, + elicitation: Option, +} + +pub type ServerFactoryError = Box; + +#[derive(Debug, Error)] +pub enum McpSpawnError { + #[error(transparent)] + Serve(#[from] ServeError), + #[error("No factory is registered for in-memory MCP server '{0}'")] + InMemoryFactoryNotFound(String), + #[error("Failed to create in-memory MCP server '{name}': {source}")] + ServerFactory { name: String, source: ServerFactoryError }, } impl McpBuilder { @@ -200,45 +56,48 @@ impl McpBuilder { Self { servers: Vec::new(), factories: HashMap::new(), - mcp_channel_capacity: 1000, root_dir: root_dir.as_ref().to_path_buf(), - oauth_handler_factory: None, agent_deps: AgentDeps::default(), - aether_home: None, vars, tool_filter: ToolFilter::default(), - progressive_discovery_instructions: None, + elicitation: None, } } - pub fn with_servers(mut self, servers: Vec) -> Self { - self.servers.extend(servers); - self + pub fn with_config(mut self, config: McpConfig) -> Result { + for (name, server) in config.servers { + let server = ConfiguredServer::from_config(name, server, &self.vars, &self.root_dir)?; + self.insert(server); + } + Ok(self) } - pub fn with_tool_filter(mut self, filter: ToolFilter) -> Self { - self.tool_filter = filter; - self + pub fn with_servers(self, servers: Vec) -> Self { + servers.into_iter().fold(self, |builder, spec| builder.with_server(spec, ToolExposure::ModelVisible)) } - pub fn with_progressive_discovery_instructions(mut self, instructions: impl Into) -> Self { - self.progressive_discovery_instructions = Some(instructions.into()); + pub fn with_server(mut self, spec: ServerSpec, exposure: ToolExposure) -> Self { + self.insert(ConfiguredServer { name: spec.name, exposure, source: ServerSource::Ready(spec.transport) }); self } - pub fn register_in_memory_server(mut self, name: impl Into, factory: ServerFactory) -> Self { - self.factories.insert(name.into(), factory); + pub fn with_tool_filter(mut self, filter: ToolFilter) -> Self { + self.tool_filter = filter; self } - pub fn root_dir(&self) -> &Path { - &self.root_dir + pub fn with_elicitation(mut self, sink: ElicitationSink) -> Self { + self.elicitation = Some(sink); + self } - /// Cross-cutting dependencies handed to every agent spawned behind this - /// builder's in-memory servers. - pub fn agent_deps(&self) -> AgentDeps { - self.agent_deps.clone() + pub fn register_in_memory_server( + mut self, + name: impl Into, + factory: impl Fn(Vec, &RuntimeServices) -> Result + Send + Sync + 'static, + ) -> Self { + self.factories.insert(name.into(), Box::new(factory)); + self } pub fn with_agent_deps(mut self, deps: AgentDeps) -> Self { @@ -246,158 +105,91 @@ impl McpBuilder { self } - pub fn with_oauth_handler_factory(mut self, factory: OAuthHandlerFactory) -> Self { - self.oauth_handler_factory = Some(factory); - self - } + pub fn spawn(self) -> Result { + let exposures: HashMap = + self.servers.iter().map(|server| (server.name.clone(), server.exposure.clone())).collect(); + let defers_tools = exposures.values().any(ToolExposure::defers_tools); + let socket = defers_tools.then(UnixSocket::bind).transpose()?; + let options = self.client_options(); + let services = RuntimeServices { + root_dir: self.root_dir, + agent_deps: self.agent_deps, + deferred_tools_socket: socket.as_ref().map(|socket| socket.path().to_path_buf()), + }; + let servers = self + .servers + .into_iter() + .map(|server| server.build(&self.factories, &services)) + .collect::, _>>()?; - pub fn with_aether_home(mut self, aether_home: impl Into) -> Self { - let aether_home = aether_home.into(); - self.vars.insert("AETHER_HOME", aether_home.to_string_lossy().into_owned()); - self.aether_home = Some(aether_home); - self + Ok(McpCatalog::new(McpGateway::new(servers, &options), self.tool_filter, exposures, socket)) } - pub fn from_json_files>(mut self, paths: &[T]) -> Result { - if paths.is_empty() { - return Ok(self); + fn insert(&mut self, server: ConfiguredServer) { + match self.servers.iter_mut().find(|existing| existing.name == server.name) { + Some(existing) => *existing = server, + None => self.servers.push(server), } - let raw = McpConfig::from_json_files(paths)?; - self.servers.extend(raw.into_servers(&self.vars)?); - Ok(self) } - pub fn from_mcp_config_sources(mut self, sources: &[McpConfigSource]) -> Result { - if sources.is_empty() { - return Ok(self); + fn client_options(&self) -> ClientOptions { + let mut options = + ClientOptions::default().implementation(Implementation::new("aether", env!("CARGO_PKG_VERSION"))); + if let Some(store) = self.agent_deps.oauth_credential_store.clone() { + options = options.oauth_store(store); } - - let mut merged = McpConfig::default(); - for source in sources { - let config = match source { - McpConfigSource::File { path, defer_tools } => { - let mut config = McpConfig::from_json_file(path)?; - if *defer_tools { - config.defer_all_tools(); - } - config - } - McpConfigSource::Json(json) => McpConfig::from_json(json)?, - McpConfigSource::Inline(config) => config.clone(), - }; - merged.servers.extend(config.servers); + if let Some(sink) = self.elicitation.clone() { + options = options.elicitation(sink); } - - self.servers.extend(merged.into_servers(&self.vars)?); - Ok(self) + options } +} - pub async fn spawn(self) -> Result { - let McpBuilder { - servers, - factories, - mcp_channel_capacity, - root_dir, - oauth_handler_factory, - agent_deps, - aether_home: _, - vars: _, - tool_filter, - progressive_discovery_instructions, - } = self; - if servers.iter().any(|server| server.tool_exposure.has_deferred_tools()) - && servers.iter().any(|server| server.name == PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME) - { - return Err(McpError::ReservedServerName(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME.to_string())); - } - let (manager_tx, manager_rx) = mpsc::channel::(mcp_channel_capacity); - let (snapshot_tx, snapshot_rx) = watch::channel(Arc::new(mcp_utils::client::McpSnapshot::default())); - let (event_tx, event_rx) = mpsc::channel::(mcp_channel_capacity); - let mcp = McpHandle::new(manager_tx, snapshot_rx); - let gateway_transport = if servers.iter().any(|server| server.tool_exposure.has_deferred_tools()) { - let path = UnixSocketPath::new().map_err(|error| McpError::TransportError(error.to_string()))?; - Some(UnixSocketMcpTransport::bind(path).map_err(|error| McpError::TransportError(error.to_string()))?) - } else { - None - }; - let shell_environment = gateway_transport - .as_ref() - .map(|transport| { - BTreeMap::from([(AETHER_MCP_IPC_SOCKET.to_string(), transport.path().to_string_lossy().into_owned())]) - }) - .unwrap_or_default(); - let services = RuntimeServices { mcp: mcp.clone(), root_dir: root_dir.clone(), agent_deps, shell_environment }; - let servers = resolve_servers(servers, &factories, &services).await?; - - let mut mcp_manager = McpManager::new(event_tx, oauth_handler_factory) - .with_tool_filter(tool_filter) - .with_snapshot_sender(snapshot_tx); - if let Some(capabilities) = services.agent_deps.mcp_client_capabilities.clone() { - mcp_manager = mcp_manager.with_client_capabilities(capabilities); - } - if let Some(instructions) = progressive_discovery_instructions { - mcp_manager = mcp_manager.with_progressive_discovery_instructions(instructions); - } - if let Some(store) = services.agent_deps.oauth_credential_store.clone() { - mcp_manager = mcp_manager.with_oauth_credential_store(store); - } - mcp_manager = mcp_manager.with_root_dir(root_dir); - let pending = mcp_manager.register_pending(servers).await?; - let task = tokio::spawn(run_mcp_task(mcp_manager, manager_rx, pending)); - let gateway = gateway_transport.map(|transport| transport.spawn(GatewayService::new(mcp.clone()))); +type ServerFactory = Box, &RuntimeServices) -> Result + Send + Sync>; - Ok(McpSession { runtime: McpRuntime { mcp, handle: task, agent_sync_handle: None, gateway }, event_rx }) - } +struct ConfiguredServer { + name: String, + exposure: ToolExposure, + source: ServerSource, } -async fn resolve_servers( - servers: Vec, - factories: &HashMap, - services: &RuntimeServices, -) -> Result, McpError> { - let mut resolved = Vec::with_capacity(servers.len()); - for McpServer { name, transport, tool_exposure } in servers { - let transport = match transport { - McpTransport::Stdio { command, args, env } => RuntimeMcpTransport::Stdio { command, args, env }, - McpTransport::Http(config) => RuntimeMcpTransport::Http(config), - McpTransport::InMemory { spec } => { - let factory = factories.get(&spec.factory).ok_or_else(|| McpError::InMemoryFactoryNotFound { - server: name.clone(), - factory: spec.factory.clone(), - })?; - RuntimeMcpTransport::InMemory { server: factory(spec, services.clone()).await } +impl ConfiguredServer { + fn from_config(name: String, config: McpServerConfig, vars: &Vars, root_dir: &Path) -> Result { + let exposure = config.defer_tools().clone(); + Ok(Self { name, exposure, source: config.resolve(vars, root_dir)? }) + } + + fn build( + self, + factories: &HashMap, + services: &RuntimeServices, + ) -> Result { + let transport = match self.source { + ServerSource::Ready(transport) => transport, + ServerSource::InMemory(args) => { + let factory = factories + .get(&self.name) + .ok_or_else(|| McpSpawnError::InMemoryFactoryNotFound(self.name.clone()))?; + let server = factory(args, services) + .map_err(|source| McpSpawnError::ServerFactory { name: self.name.clone(), source })?; + Transport::InProcess(server) } }; - resolved.push(RuntimeMcpServer::new(name, transport, tool_exposure)); + Ok(ServerSpec { name: self.name, transport }) } - Ok(resolved) } #[cfg(test)] mod tests { use super::*; use aether_auth::{FakeOAuthCredentialStore, OAuthCredentialStorage}; - use futures::FutureExt; - use mcp_utils::client::{McpServerConfig, McpTransport, StdioServerConfig, StdioType, ToolExposure}; use mcp_utils::testing::FakeMcpServer; - use std::collections::{BTreeMap, HashMap}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use utils::mcp_status::McpServerStatus; - fn write_config_file(name: &str, json: &str) -> (tempfile::TempDir, PathBuf) { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join(name); - std::fs::write(&path, json).unwrap(); - (dir, path) - } - - fn json_source(json: &str) -> McpConfigSource { - McpConfigSource::Json(json.to_string()) - } - - fn builder_from_sources(sources: &[McpConfigSource]) -> McpBuilder { - McpBuilder::new("/workspace").from_mcp_config_sources(sources).unwrap() + fn config(json: &str) -> McpConfig { + McpConfig::from_json(json).unwrap() } #[tokio::test] @@ -408,257 +200,169 @@ mod tests { let factory_received = Arc::clone(&received); let oauth_store: Arc = Arc::new(FakeOAuthCredentialStore::new()); let deps = AgentDeps::new(Arc::clone(&oauth_store), None); - let factory: ServerFactory = Box::new(move |spec, services| { - factory_calls.fetch_add(1, Ordering::SeqCst); - assert_eq!(spec.args, ["--root", "/workspace/tools"]); - assert_eq!(spec.input, Some(serde_json::json!({"enabled": true}))); - *factory_received.lock().unwrap() = Some(services); - async move { FakeMcpServer::new().into_dyn() }.boxed() - }); let builder = McpBuilder::new("/workspace") .with_agent_deps(deps) - .register_in_memory_server("test", factory) - .from_mcp_config_sources(&[json_source( - r#"{"servers":{"test":{"type":"in-memory","args":["--root","${WORKSPACE}/tools"],"input":{"enabled":true}}}}"#, - )]) + .register_in_memory_server("test", move |args, services| { + factory_calls.fetch_add(1, Ordering::SeqCst); + assert_eq!(args, ["--root", "/workspace/tools"]); + *factory_received.lock().unwrap() = Some(services.clone()); + Ok(FakeMcpServer::new().into()) + }) + .with_config(config(r#"{"servers":{"test":{"type":"in-memory","args":["--root","${WORKSPACE}/tools"]}}}"#)) .unwrap(); assert_eq!(calls.load(Ordering::SeqCst), 0); - let spawn = builder.spawn().await.unwrap(); + let _spawned = builder.spawn().unwrap(); assert_eq!(calls.load(Ordering::SeqCst), 1); let services = received.lock().unwrap().clone().expect("factory received runtime services"); assert_eq!(services.root_dir, PathBuf::from("/workspace")); - assert!(Arc::ptr_eq(&services.mcp.snapshot(), &spawn.handle().snapshot())); assert!(Arc::ptr_eq( services.agent_deps.oauth_credential_store.as_ref().expect("factory received agent dependencies"), &oauth_store, )); - assert!(services.shell_environment.is_empty()); + assert!(services.deferred_tools_socket.is_none()); } #[tokio::test] async fn deferred_gateway_is_bound_before_in_memory_factories_run() { let received = Arc::new(Mutex::new(None::)); let factory_received = Arc::clone(&received); - let factory: ServerFactory = Box::new(move |_, services| { - *factory_received.lock().unwrap() = Some(services); - async move { FakeMcpServer::new().into_dyn() }.boxed() - }); - let spawn = McpBuilder::new("/workspace") - .register_in_memory_server("test", factory) - .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory","deferTools":true}}}"#)]) + let mcp = McpBuilder::new("/workspace") + .register_in_memory_server("test", move |_, services| { + *factory_received.lock().unwrap() = Some(services.clone()); + Ok(FakeMcpServer::new().into()) + }) + .with_config(config(r#"{"servers":{"test":{"type":"in-memory","deferTools":true}}}"#)) .unwrap() .spawn() - .await .unwrap(); let services = received.lock().unwrap().clone().expect("factory received runtime services"); - let inherited = - services.shell_environment.get(AETHER_MCP_IPC_SOCKET).expect("factory receives gateway endpoint"); - assert_eq!(Path::new(inherited), spawn.gateway_endpoint().expect("gateway endpoint exists")); - assert!(Path::new(inherited).exists()); + let socket = services.deferred_tools_socket.expect("factory receives the deferred tools socket"); + assert_eq!(socket, mcp.deferred_tools_socket().expect("the session serves deferred tools")); + assert!(socket.exists()); } #[tokio::test] - async fn snapshots_are_immutable_and_watch_observes_connection_changes() { - let factory: ServerFactory = Box::new(|_, _| async move { FakeMcpServer::new().into_dyn() }.boxed()); - let mut spawn = McpBuilder::new("/workspace") - .register_in_memory_server("test", factory) - .from_mcp_config_sources(&[json_source(r#"{"servers":{"test":{"type":"in-memory"}}}"#)]) - .unwrap() - .spawn() - .await - .unwrap(); - let old = spawn.handle().snapshot(); - let mut updates = spawn.handle().subscribe(); + async fn shutting_down_a_session_stops_the_deferred_socket_and_the_servers_for_every_clone() { + let lazy = ServerSpec::in_process("lazy", FakeMcpServer::new()); + let mcp = McpBuilder::new("/workspace").with_server(lazy, ToolExposure::deferred_all()).spawn().unwrap(); + mcp.ready().await; + let held = mcp.clone(); + let socket = mcp.deferred_tools_socket().expect("the session serves deferred tools").to_path_buf(); - let ready = spawn.block_until_ready().await.expect("bootstrap completes"); - updates.changed().await.expect("connection publishes a snapshot"); - let observed = updates.borrow().clone(); + mcp.shutdown().await; - assert_eq!(old.tool_definitions().len(), 0); - assert_eq!(ready.tool_definitions()[0].name, "test__add_numbers"); - assert_eq!(observed.tool_definitions(), ready.tool_definitions()); - assert!(!Arc::ptr_eq(&old, &ready)); + assert!(!socket.exists()); + assert_eq!(held.statuses().len(), 0); } #[tokio::test] - async fn missing_in_memory_factory_fails_at_spawn_with_server_and_factory() { - let builder = McpBuilder::new("/workspace") - .from_mcp_config_sources(&[json_source(r#"{"servers":{"custom":{"type":"in-memory"}}}"#)]) + async fn spawned_in_memory_servers_connect_in_the_background() { + let mcp = McpBuilder::new("/workspace") + .register_in_memory_server("test", |_, _| Ok(FakeMcpServer::math().into())) + .with_config(config(r#"{"servers":{"test":{"type":"in-memory"}}}"#)) + .unwrap() + .spawn() .unwrap(); - let Err(error) = builder.spawn().await else { - panic!("spawn should reject an unregistered factory"); - }; - assert!(matches!( - error, - McpError::InMemoryFactoryNotFound { ref server, ref factory } - if server == "custom" && factory == "custom" - )); - } + mcp.ready().await; - #[tokio::test] - async fn mixed_direct_sources_preserve_last_wins_order() { - let (_dir, file_path) = - write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#); - let inline = McpConfig::new(BTreeMap::from([( - "coding".to_string(), - McpServerConfig::Stdio(StdioServerConfig { - type_: StdioType::Stdio, - command: "from_inline".to_string(), - args: Vec::new(), - env: HashMap::new(), - defer_tools: ToolExposure::ModelVisible, - }), - )])); - let sources = vec![ - McpConfigSource::model_visible(file_path), - json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#), - McpConfigSource::Inline(inline), - ]; - - let builder = builder_from_sources(&sources); - - assert_eq!(command_for(&builder, "coding"), Some("from_inline")); - assert_eq!(deferred_tools_for(&builder, "coding"), Some(false)); + assert_eq!(mcp.tool_definitions()[0].name, "test__add_numbers"); } #[tokio::test] - async fn file_sources_keep_their_position_relative_to_json_sources() { - let (_dir, file_path) = - write_config_file("mcp.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#); - let sources = vec![ - json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json"}}}"#), - McpConfigSource::model_visible(file_path), - ]; + async fn ready_servers_and_config_servers_are_added_in_order() { + let mcp = McpBuilder::new("/workspace") + .register_in_memory_server("configured", |_, _| Ok(FakeMcpServer::new().into())) + .with_servers(vec![ServerSpec::in_process("ready", FakeMcpServer::new())]) + .with_config(config(r#"{"servers":{"configured":{"type":"in-memory"}}}"#)) + .unwrap() + .spawn() + .unwrap(); - let builder = builder_from_sources(&sources); + mcp.ready().await; - assert_eq!(command_for(&builder, "coding"), Some("from_file")); + assert_eq!(names(&mcp), ["ready", "configured"]); } #[tokio::test] - async fn file_source_defer_tools_marks_all_file_servers_deferred() { - let (_dir, file_path) = write_config_file( - "deferred.json", - r#"{"servers":{"github":{"type":"stdio","command":"g","deferTools":{"exclude":["status"]}},"browser":{"type":"stdio","command":"b"}}}"#, - ); - - let builder = McpBuilder::new("/workspace") - .from_mcp_config_sources(&[McpConfigSource::File { path: file_path, defer_tools: true }]) + async fn a_server_replaces_an_earlier_one_with_the_same_name_in_its_position() { + let replaced = FakeMcpServer::new(); + let alive = replaced.alive(); + let replacement = FakeMcpServer::new().with_tool(mcp_utils::testing::FakeTool::new("replacement")); + let mcp = McpBuilder::new("/workspace") + .with_servers(vec![ + ServerSpec::in_process("math", replaced), + ServerSpec::in_process("docs", FakeMcpServer::new()), + ]) + .with_server(ServerSpec::in_process("math", replacement), ToolExposure::deferred_all()) + .spawn() .unwrap(); - assert_eq!(deferred_tools_for(&builder, "github"), Some(true)); - assert_eq!(deferred_tools_for(&builder, "browser"), Some(true)); - assert!(is_direct_tool(&builder, "github", "status")); - } + mcp.ready().await; - #[tokio::test] - async fn later_sources_override_defer_tools_flag() { - let (_dir, file_path) = - write_config_file("deferred.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_file"}}}"#); - let sources = vec![ - McpConfigSource::File { path: file_path, defer_tools: true }, - json_source(r#"{"servers":{"coding":{"type":"stdio","command":"from_json","deferTools":false}}}"#), - ]; - - let builder = builder_from_sources(&sources); - - assert_eq!(command_for(&builder, "coding"), Some("from_json")); - assert_eq!(deferred_tools_for(&builder, "coding"), Some(false)); + assert!(alive.upgrade().is_none()); + assert_eq!(names(&mcp), ["math", "docs"]); + assert!(mcp.statuses()[0].deferred_tools); + assert_eq!(mcp.statuses()[0].status, McpServerStatus::Connected { tool_count: 1 }); } #[tokio::test] - async fn spawn_returns_immediately_and_emits_initial_connecting_status() { - let spawn = McpBuilder::new("/workspace") - .from_mcp_config_sources(&[json_source( - r#"{"servers":{"slow":{"type":"stdio","command":"sleep","args":["30"]}}}"#, - )]) - .unwrap() - .spawn() - .await - .expect("spawn should succeed"); + async fn missing_in_memory_factory_fails_at_spawn_with_server_name() { + let builder = McpBuilder::new("/workspace") + .with_config(config(r#"{"servers":{"custom":{"type":"in-memory"}}}"#)) + .unwrap(); - let (_runtime, mut event_rx) = spawn.split(); - let event = event_rx.try_recv().expect("spawn() should buffer an initial ServerStatusesChanged"); - let McpClientEvent::ServerStatusesChanged(statuses) = event else { - panic!("expected ServerStatusesChanged, got {event:?}"); + let Err(error) = builder.spawn() else { + panic!("spawn should reject an unregistered factory"); }; - assert!(matches!(statuses[0].status, McpServerStatus::Connecting)); + assert!(matches!(error, McpSpawnError::InMemoryFactoryNotFound(ref server) if server == "custom")); } #[tokio::test] - async fn from_mcp_config_sources_expands_workspace_var_in_stdio_args() { - let builder = McpBuilder::new("/work") - .from_mcp_config_sources(&[json_source( - r#"{"servers":{"notes":{"type":"stdio","command":"server","args":["--dir","${WORKSPACE}/notes"]}}}"#, - )]) + async fn rejected_in_memory_args_fail_spawn_with_server_name() { + let builder = McpBuilder::new("/workspace") + .register_in_memory_server("custom", |args, _| Err(format!("unexpected argument '{}'", args[0]).into())) + .with_config(config(r#"{"servers":{"custom":{"type":"in-memory","args":["--bad"]}}}"#)) .unwrap(); - assert_eq!(args_for(&builder, "notes"), Some(vec!["--dir".to_string(), "/work/notes".to_string()])); + let Err(error) = builder.spawn() else { + panic!("spawn should surface the factory's error"); + }; + assert!(matches!(error, McpSpawnError::ServerFactory { ref name, .. } if name == "custom")); + assert!(error.to_string().contains("unexpected argument '--bad'"), "unexpected error: {error}"); } #[tokio::test] - async fn from_mcp_config_sources_expands_aether_home_var_in_stdio_args() { - let home = tempfile::tempdir().unwrap(); - - let builder = McpBuilder::new("/work") - .with_aether_home(home.path()) - .from_mcp_config_sources(&[json_source( - r#"{"servers":{"skills":{"type":"stdio","command":"server","args":["--dir","${AETHER_HOME}/skills"]}}}"#, - )]) - .unwrap(); + async fn config_args_expand_the_aether_home_var() { + let home = SettingsStore::new("AETHER_HOME", ".aether").expect("Aether home resolves").home().to_path_buf(); - assert_eq!( - args_for(&builder, "skills"), - Some(vec!["--dir".to_string(), home.path().join("skills").to_string_lossy().into_owned()]) + let args = spawned_args( + "/work", + r#"{"servers":{"notes":{"type":"in-memory","args":["--dir","${AETHER_HOME}/skills"]}}}"#, ); - } - #[tokio::test] - async fn reserved_progressive_discovery_server_is_rejected_when_gateway_is_enabled() { - let result = McpBuilder::new("/workspace") - .from_mcp_config_sources(&[json_source( - r#"{"servers":{"progressive-discovery":{"type":"stdio","command":"server"},"deferred":{"type":"stdio","command":"server","deferTools":true}}}"#, - )]) - .unwrap() - .spawn() - .await; - - assert!(matches!(result, Err(McpError::ReservedServerName(name)) if name == "progressive-discovery")); - } - - #[test] - fn new_sets_root_directory_from_workspace_root() { - let builder = McpBuilder::new("/workspace"); - assert_eq!(builder.root_dir, PathBuf::from("/workspace")); - } - - fn command_for<'a>(builder: &'a McpBuilder, name: &str) -> Option<&'a str> { - builder.servers.iter().find_map(|server| match &server.transport { - McpTransport::Stdio { command, .. } if server.name == name => Some(command.as_str()), - _ => None, - }) - } - - fn args_for(builder: &McpBuilder, name: &str) -> Option> { - builder.servers.iter().find_map(|server| match &server.transport { - McpTransport::Stdio { args, .. } if server.name == name => Some(args.clone()), - _ => None, - }) + assert_eq!(args, ["--dir".to_string(), home.join("skills").to_string_lossy().into_owned()]); } - fn is_direct_tool(builder: &McpBuilder, server_name: &str, tool_name: &str) -> bool { - builder - .servers - .iter() - .find(|server| server.name == server_name) - .is_some_and(|server| server.tool_exposure.is_model_visible_tool(tool_name)) + fn names(mcp: &McpCatalog) -> Vec { + mcp.statuses().into_iter().map(|status| status.name).collect() } - fn deferred_tools_for(builder: &McpBuilder, name: &str) -> Option { - builder.servers.iter().find(|server| server.name == name).map(mcp_utils::client::McpServer::has_deferred_tools) + fn spawned_args(root_dir: &str, json: &str) -> Vec { + let received = Arc::new(Mutex::new(Vec::new())); + let factory_received = Arc::clone(&received); + McpBuilder::new(root_dir) + .register_in_memory_server("notes", move |args, _| { + *factory_received.lock().unwrap() = args; + Ok(FakeMcpServer::new().into()) + }) + .with_config(config(json)) + .unwrap() + .spawn() + .unwrap(); + received.lock().unwrap().clone() } } diff --git a/crates/aether-core/src/mcp/mcp_catalog.rs b/crates/aether-core/src/mcp/mcp_catalog.rs new file mode 100644 index 000000000..290a6d525 --- /dev/null +++ b/crates/aether-core/src/mcp/mcp_catalog.rs @@ -0,0 +1,202 @@ +use super::deferred_tools; +use super::tool_bridge::tool_definition; +use super::{ToolExposure, ToolFilter}; +use crate::core::McpInstructions; +use futures::Stream; +use llm::ToolDefinition; +use mcp_utils::client::{ClientOptions, ToolCall, ToolCallOptions}; +use mcp_utils::gateway::{ConnectedServer, GatewayError, McpGateway, ServerPrompt, ServerState, ServerStatus}; +use mcp_utils::model::{GetPromptResult, Tool}; +use mcp_utils::server::{ServerHandle, UnixSocket}; +use serde_json::{Map, Value}; +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; +use utils::mcp_status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry}; + +#[derive(Clone)] +pub struct McpCatalog { + core: CatalogCore, + deferred_tools_server: Option>, +} + +impl McpCatalog { + pub async fn ready(&self) { + self.core.gateway.ready().await; + } + + pub fn changes(&self) -> impl Stream + Send + Sync + Unpin + use<> { + self.core.gateway.changes() + } + + pub fn statuses(&self) -> Vec { + self.core.gateway.statuses().into_iter().map(|status| self.core.policy.status_entry(status)).collect() + } + + pub fn tool_definitions(&self) -> Vec { + let tools = self.core.tools(Tier::Direct).into_iter(); + tools.map(|(server, tool)| tool_definition(&tool).with_server(server)).collect() + } + + pub fn instructions(&self) -> McpInstructions { + let servers = self.core.servers(Tier::Direct).into_iter(); + McpInstructions { + servers: servers + .filter_map(|(server, _)| Some((server.name().to_string(), server.instructions()?))) + .collect(), + progressive_discovery: !self.core.servers(Tier::Deferred).is_empty(), + } + } + + pub fn call_tool(&self, name: &str, args: Map, options: ToolCallOptions) -> Option { + self.core.call(Tier::Direct, name, args, options) + } + + pub async fn list_prompts(&self) -> Vec { + self.core.gateway.list_prompts().await + } + + pub async fn get_prompt( + &self, + server: &str, + prompt: &str, + args: Map, + ) -> Result { + self.core.gateway.get_prompt(server, prompt, args).await + } + + pub async fn authenticate(&self, server: &str) -> Result<(), GatewayError> { + self.core.gateway.authenticate(server).await + } + + pub fn deferred_tools_socket(&self) -> Option<&Path> { + self.deferred_tools_server.as_deref().map(ServerHandle::path) + } + + pub async fn shutdown(&self) { + if let Some(server) = &self.deferred_tools_server { + server.stop(); + } + self.core.gateway.shutdown().await; + } + + pub(crate) fn new( + gateway: McpGateway, + filter: ToolFilter, + exposures: HashMap, + socket: Option, + ) -> Self { + let core = CatalogCore { gateway, policy: Arc::new(ToolPolicy { filter, exposures }) }; + let deferred_tools_server = socket.map(|socket| Arc::new(deferred_tools::serve(core.clone(), socket))); + Self { core, deferred_tools_server } + } +} + +impl Default for McpCatalog { + fn default() -> Self { + let gateway = McpGateway::new(Vec::new(), &ClientOptions::default()); + Self::new(gateway, ToolFilter::default(), HashMap::new(), None) + } +} + +pub(super) fn namespaced(server: &str, tool: &str) -> String { + format!("{server}{SERVER_DELIMITER}{tool}") +} + +const SERVER_DELIMITER: &str = "__"; + +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum Tier { + Direct, + Deferred, +} + +#[derive(Clone)] +pub(super) struct CatalogCore { + gateway: McpGateway, + policy: Arc, +} + +impl CatalogCore { + pub(super) fn servers(&self, tier: Tier) -> Vec<(ConnectedServer, Vec)> { + let servers = self.gateway.servers().into_iter(); + servers + .filter_map(|server| { + let offered = server.tools().iter().filter(|tool| self.policy.tier(server.name(), tool) == Some(tier)); + let offered = offered.cloned().collect::>(); + (!offered.is_empty()).then_some((server, offered)) + }) + .collect() + } + + pub(super) fn tools(&self, tier: Tier) -> Vec<(String, Tool)> { + let servers = self.servers(tier).into_iter(); + servers + .flat_map(|(server, tools)| { + let server = server.name().to_string(); + tools.into_iter().map(move |mut tool| { + tool.name = namespaced(&server, &tool.name).into(); + (server.clone(), tool) + }) + }) + .collect() + } + + pub(super) fn tier_of(&self, name: &str) -> Option { + self.resolve(name).map(|(_, _, tier)| tier) + } + + pub(super) fn call( + &self, + tier: Tier, + name: &str, + args: Map, + options: ToolCallOptions, + ) -> Option { + let (server, tool, offered_in) = self.resolve(name)?; + (offered_in == tier).then(|| server.call_tool(tool, args, options)) + } + + fn resolve<'n>(&self, name: &'n str) -> Option<(ConnectedServer, &'n str, Tier)> { + let (server, tool) = self.split(name)?; + let server = self.gateway.server(server)?; + let tier = self.policy.tier(server.name(), server.tools().iter().find(|offered| offered.name == tool)?)?; + Some((server, tool, tier)) + } + + fn split<'n>(&self, name: &'n str) -> Option<(&str, &'n str)> { + let servers = self.policy.exposures.keys(); + servers + .filter_map(|server| Some((server.as_str(), name.strip_prefix(server)?.strip_prefix(SERVER_DELIMITER)?))) + .max_by_key(|(server, _)| server.len()) + } +} + +struct ToolPolicy { + filter: ToolFilter, + exposures: HashMap, +} + +impl ToolPolicy { + fn tier(&self, server: &str, tool: &Tool) -> Option { + if !self.filter.allows(&namespaced(server, &tool.name), tool) { + return None; + } + let deferred = self.exposures.get(server).is_some_and(|exposure| exposure.is_deferred(tool)); + Some(if deferred { Tier::Deferred } else { Tier::Direct }) + } + + fn status_entry(&self, status: ServerStatus) -> McpServerStatusEntry { + let defers_tools = self.exposures.get(&status.name).is_some_and(ToolExposure::defers_tools); + let auth = + if status.can_sign_in { McpServerAuthCapability::OAuth } else { McpServerAuthCapability::Unavailable }; + let state = match status.state { + ServerState::Connecting => McpServerStatus::Connecting, + ServerState::Authenticating => McpServerStatus::Authenticating, + ServerState::NeedsOAuth => McpServerStatus::NeedsOAuth, + ServerState::Failed { error } => McpServerStatus::Failed { error }, + ServerState::Connected { tool_count } => McpServerStatus::Connected { tool_count }, + }; + McpServerStatusEntry::new(status.name, state).with_auth_capability(auth).with_deferred_tools(defers_tools) + } +} diff --git a/crates/aether-core/src/mcp/mcp_handle.rs b/crates/aether-core/src/mcp/mcp_handle.rs deleted file mode 100644 index 981216c3b..000000000 --- a/crates/aether-core/src/mcp/mcp_handle.rs +++ /dev/null @@ -1,132 +0,0 @@ -use super::run_mcp_task::ManagerCommand; -use futures::{Stream, future::join_all}; -use mcp_utils::client::{CallToolError, CallToolOptions, McpError, McpSnapshot, ToolCallEvent, ToolRoute, call_tool}; -use rmcp::model::{GetPromptRequestParams, GetPromptResult, Prompt}; -use serde_json::{Map, Value}; -use std::{pin::Pin, sync::Arc}; -use thiserror::Error; -use tokio::sync::{mpsc, oneshot, watch}; - -pub type ToolCallStream = Pin + Send>>; - -#[derive(Clone)] -pub struct McpHandle { - control_tx: mpsc::Sender, - snapshot_rx: watch::Receiver>, -} - -#[derive(Debug, Error)] -pub enum McpHandleError { - #[error(transparent)] - Route(#[from] McpError), - #[error("MCP manager is unavailable")] - ManagerUnavailable, - #[error("failed to list prompts for {server}: {message}")] - PromptList { server: String, message: String }, - #[error("failed to get prompt '{prompt}' from {server}: {message}")] - PromptGet { server: String, prompt: String, message: String }, -} - -impl McpHandle { - pub(super) fn new( - control_tx: mpsc::Sender, - snapshot_rx: watch::Receiver>, - ) -> Self { - Self { control_tx, snapshot_rx } - } - - pub fn snapshot(&self) -> Arc { - self.snapshot_rx.borrow().clone() - } - - pub fn subscribe(&self) -> watch::Receiver> { - self.snapshot_rx.clone() - } - - pub fn call(&self, route: ToolRoute, arguments: Map, options: CallToolOptions) -> ToolCallStream { - match self.snapshot().resolve(route, arguments) { - Ok((client, params)) => Box::pin(call_tool(client, params, options)), - Err(error) => Box::pin(futures::stream::once(async move { - ToolCallEvent::Complete(Err(CallToolError::Unavailable { - message: format!("Failed to resolve tool: {error}"), - })) - })), - } - } - - pub fn call_model_visible( - &self, - namespaced_name: String, - arguments_json: &str, - options: CallToolOptions, - ) -> ToolCallStream { - let parsed = serde_json::from_str::(arguments_json) - .map_err(McpError::from) - .and_then(|value| { - value - .as_object() - .cloned() - .ok_or_else(|| McpError::JsonError("tool arguments must be a JSON object".to_string())) - }) - .map(|arguments| (ToolRoute::ModelVisible { namespaced_name }, arguments)); - match parsed { - Ok((route, arguments)) => self.call(route, arguments, options), - Err(error) => failed_call(error), - } - } - - pub async fn list_prompts(&self) -> Result, McpHandleError> { - let futures = self.snapshot().clients_with_prompts().into_iter().map(|(server, client)| async move { - let prompts = client - .list_all_prompts() - .await - .map_err(|error| McpHandleError::PromptList { server: server.clone(), message: error.to_string() })?; - Ok::<_, McpHandleError>( - prompts - .into_iter() - .map(|prompt| { - Prompt::new(format!("{server}__{}", prompt.name), prompt.description, prompt.arguments) - }) - .collect::>(), - ) - }); - let mut prompts = Vec::new(); - for result in join_all(futures).await { - prompts.extend(result?); - } - Ok(prompts) - } - - pub async fn get_prompt( - &self, - name: &str, - arguments: Option>, - ) -> Result { - let (prompt, client) = self.snapshot().client_for_prompt(name)?; - let server = client.service().server_name().to_string(); - let mut request = GetPromptRequestParams::new(prompt.clone()); - if let Some(arguments) = arguments { - request = request.with_arguments(arguments); - } - client.get_prompt(request).await.map_err(|error| McpHandleError::PromptGet { - server, - prompt, - message: error.to_string(), - }) - } - - pub async fn authenticate_server(&self, name: &str) -> Result<(), McpHandleError> { - let (tx, rx) = oneshot::channel(); - self.control_tx - .send(ManagerCommand::AuthenticateServer { name: name.to_string(), tx }) - .await - .map_err(|_| McpHandleError::ManagerUnavailable)?; - rx.await.map_err(|_| McpHandleError::ManagerUnavailable)?.map_err(McpHandleError::Route) - } -} - -fn failed_call(error: McpError) -> ToolCallStream { - Box::pin(futures::stream::once(async move { - ToolCallEvent::Complete(Err(CallToolError::Unavailable { message: format!("Failed to resolve tool: {error}") })) - })) -} diff --git a/crates/aether-core/src/mcp/mod.rs b/crates/aether-core/src/mcp/mod.rs index c032dd8b1..ef0d6cf8b 100644 --- a/crates/aether-core/src/mcp/mod.rs +++ b/crates/aether-core/src/mcp/mod.rs @@ -1,10 +1,20 @@ -pub mod mcp_builder; -pub mod tool_bridge; +mod config; +mod deferred_tools; +mod exposure; +mod mcp_builder; +mod mcp_catalog; +pub(crate) mod tool_bridge; +mod tool_filter; -mod gateway_service; -mod mcp_handle; -mod run_mcp_task; - -pub use gateway_service::GatewayService; -pub use mcp_builder::*; -pub use mcp_handle::{McpHandle, McpHandleError, ToolCallStream}; +pub use config::{ + InMemoryServerConfig, InMemoryType, McpConfig, McpOAuthConfig, McpServerConfig, ParseError, RemoteServerConfig, + RemoteType, StdioServerConfig, StdioType, http_transport, +}; +pub(crate) use deferred_tools::PROGRESSIVE_DISCOVERY_INSTRUCTIONS; +pub use deferred_tools::{ + AETHER_MCP_IPC_SOCKET, DeferredServer, DeferredTool, DeferredToolsClient, DeferredToolsError, +}; +pub use exposure::ToolExposure; +pub use mcp_builder::{McpBuilder, McpSpawnError, RuntimeServices, ServerFactoryError, mcp}; +pub use mcp_catalog::McpCatalog; +pub use tool_filter::{ToolAnnotationMatcher, ToolFilter, ToolMatcher}; diff --git a/crates/mcp-servers/src/progressive_discovery_instructions.md b/crates/aether-core/src/mcp/progressive_discovery_instructions.md similarity index 100% rename from crates/mcp-servers/src/progressive_discovery_instructions.md rename to crates/aether-core/src/mcp/progressive_discovery_instructions.md diff --git a/crates/aether-core/src/mcp/run_mcp_task.rs b/crates/aether-core/src/mcp/run_mcp_task.rs deleted file mode 100644 index 5ac9a0d52..000000000 --- a/crates/aether-core/src/mcp/run_mcp_task.rs +++ /dev/null @@ -1,91 +0,0 @@ -use mcp_utils::client::{McpConnectAttempt, McpConnectionAttemptManager, McpError, McpManager, RuntimeMcpServer}; -use std::collections::HashSet; -use std::time::Duration; -use tokio::select; -use tokio::sync::{mpsc, oneshot}; -use tokio::task::JoinSet; - -const MCP_AUTH_TIMEOUT: Duration = Duration::from_mins(3); - -#[derive(Debug)] -pub(super) enum ManagerCommand { - AuthenticateServer { name: String, tx: oneshot::Sender> }, -} - -pub(super) async fn run_mcp_task( - mut mcp: McpManager, - mut command_rx: mpsc::Receiver, - pending_servers: Vec, -) { - let mut tool_refresh_rx = mcp.take_tool_refresh_receiver(); - let mut tool_refreshes = JoinSet::new(); - let mut mcp_connection_attempts = McpConnectionAttemptManager::default(); - let mut pending_connections: HashSet = pending_servers.iter().map(|server| server.name.clone()).collect(); - for server in pending_servers { - let name = server.name.clone(); - let task = mcp.connect_pending_task(server); - mcp_connection_attempts.spawn(name, task); - } - if pending_connections.is_empty() { - mcp.emit_connection_ready().await; - } - - loop { - select! { - command = command_rx.recv() => { - let Some(command) = command else { break; }; - on_command(command, &mut mcp, &mut mcp_connection_attempts).await; - } - Some(joined) = mcp_connection_attempts.join_next(), if !mcp_connection_attempts.is_empty() => { - match joined { - Ok(attempt) => { - let was_bootstrap = pending_connections.remove(&attempt.name); - mcp.apply_connection_attempt(attempt).await; - if was_bootstrap && pending_connections.is_empty() { - mcp.emit_connection_ready().await; - } - } - Err(error) => tracing::error!("MCP auth task did not complete normally: {error:?}"), - } - } - Some(request) = tool_refresh_rx.recv() => { - tool_refreshes.spawn(request.refresh()); - } - Some(joined) = tool_refreshes.join_next(), if !tool_refreshes.is_empty() => { - match joined { - Ok(refresh) => mcp.apply_tool_list_refresh(refresh).await, - Err(error) => tracing::warn!(%error, "MCP tool refresh task did not complete normally"), - } - } - } - } - - mcp_connection_attempts.shutdown().await; - tool_refreshes.abort_all(); - while tool_refreshes.join_next().await.is_some() {} - mcp.shutdown().await; - tracing::debug!("MCP manager task ended"); -} - -async fn on_command(command: ManagerCommand, mcp: &mut McpManager, auth_tasks: &mut McpConnectionAttemptManager) { - match command { - ManagerCommand::AuthenticateServer { name, tx } => match mcp.authenticate_server_task(&name).await { - Ok(task) => { - let server_name = name.clone(); - auth_tasks.spawn(name, async move { - match tokio::time::timeout(MCP_AUTH_TIMEOUT, task).await { - Ok(attempt) => attempt, - Err(_) => McpConnectAttempt::failed( - server_name, - McpError::ConnectionFailed("authentication timed out after 3 minutes".to_string()), - ), - } - }); - let _ = tx.send(Ok(())); - } - Err(error) => { - let _ = tx.send(Err(error)); - } - }, - } -} diff --git a/crates/aether-core/src/mcp/tool_bridge.rs b/crates/aether-core/src/mcp/tool_bridge.rs index 2d638db5e..462dce65d 100644 --- a/crates/aether-core/src/mcp/tool_bridge.rs +++ b/crates/aether-core/src/mcp/tool_bridge.rs @@ -1,37 +1,28 @@ +use super::McpCatalog; use crate::events::{TaskOutcome, TaskOutcomeState}; -use mcp_utils::client::{CallToolError, SERVERNAME_DELIMITER}; -use rmcp::model::{CallToolRequestParams, CallToolResult, ContentBlock, EmbeddedResource, ResourceContents, Task}; +use mcp_utils::client::{ToolCall, ToolCallError as McpToolCallError, ToolCallOptions}; +use rmcp::model::{CallToolResult, ContentBlock, EmbeddedResource, ResourceContents, Tool}; +use serde_json::Value; +use thiserror::Error; -use llm::{ToolCallError, ToolCallRequest, ToolCallResult}; +use llm::{ToolAnnotations, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition}; use utils::display_meta::ToolResultMeta; use utils::temp_dir::TempDir; use utils::tool_result_truncator::ToolResultTruncator; const TOOL_RESULT_TRUNCATOR: ToolResultTruncator = ToolResultTruncator { head: 25_000, tail: 25_000 }; -/// Convert a `ToolCallRequest` to `rmcp::CallToolRequestParams` -pub fn tool_call_request_to_mcp(request: &ToolCallRequest) -> Result { - let tool_name = request - .name - .split_once(SERVERNAME_DELIMITER) - .map_or_else(|| request.name.clone(), |(_, tool_name)| tool_name.to_string()); - - // Parse arguments from JSON string - let arguments = serde_json::from_str::(&request.arguments) - .map_err(|e| format!("Invalid tool arguments: {e}"))? - .as_object() - .cloned(); - - let mut params = CallToolRequestParams::new(tool_name); - if let Some(args) = arguments { - params = params.with_arguments(args); - } - Ok(params) +#[derive(Debug, Error)] +pub(crate) enum ToolRejection { + #[error("Tool not found: {0}")] + NotFound(String), + #[error("Invalid tool arguments: {0}")] + InvalidArguments(#[source] serde_json::Error), } -pub fn convert_tool_result( +pub(crate) fn convert_tool_result( request: &ToolCallRequest, - outcome: Result, + outcome: Result, spill_dir: &TempDir, ) -> Result<(ToolCallResult, Option), ToolCallError> { let mcp_result = outcome.map_err(|error| ToolCallError::from_request(request, error.to_string()))?; @@ -61,23 +52,51 @@ pub fn convert_tool_result( )) } -pub fn map_task_result_to_outcome( +pub(crate) fn map_task_result_to_outcome( request: ToolCallRequest, - task: Task, - outcome: Result, + task_id: String, + outcome: Result, spill_dir: &TempDir, ) -> TaskOutcome { - let state = match convert_tool_result(&request, outcome, spill_dir) { - Ok((result, result_meta)) => TaskOutcomeState::Completed { result, result_meta }, - Err(error) => TaskOutcomeState::Failed { error }, + let state = match outcome { + Err(McpToolCallError::Cancelled) => TaskOutcomeState::Cancelled, + outcome => match convert_tool_result(&request, outcome, spill_dir) { + Ok((result, result_meta)) => TaskOutcomeState::Completed { result, result_meta }, + Err(error) => TaskOutcomeState::Failed { error }, + }, }; - TaskOutcome { request, task_id: task.task_id, state } + TaskOutcome { request, task_id, state } +} + +pub(crate) fn call_tool( + mcp: &McpCatalog, + request: &ToolCallRequest, + options: ToolCallOptions, +) -> Result { + let arguments = serde_json::from_str(&request.arguments).map_err(ToolRejection::InvalidArguments)?; + mcp.call_tool(&request.name, arguments, options).ok_or_else(|| ToolRejection::NotFound(request.name.clone())) } pub(crate) fn encode_structured(value: &serde_json::Value) -> String { noyalib::to_string(value).unwrap_or_else(|_| value.to_string()) } +pub(super) fn tool_definition(tool: &Tool) -> ToolDefinition { + let annotations = tool.annotations.as_ref().map(|annotations| ToolAnnotations { + title: annotations.title.clone(), + read_only_hint: annotations.read_only_hint, + destructive_hint: annotations.destructive_hint, + idempotent_hint: annotations.idempotent_hint, + open_world_hint: annotations.open_world_hint, + }); + ToolDefinition::new( + tool.name.to_string(), + tool.description.clone().unwrap_or_default(), + Value::Object((*tool.input_schema).clone()), + ) + .with_annotations(annotations) +} + fn content_text(content: &[ContentBlock], empty: &str) -> String { if content.is_empty() { return empty.to_string(); @@ -141,7 +160,7 @@ mod tests { } fn convert(mcp: McpCallToolResult) -> Result<(ToolCallResult, Option), ToolCallError> { - convert_tool_result(&req(), Ok(mcp), &TempDir::new()) + convert_tool_result(&req(), Ok::<_, McpToolCallError>(mcp), &TempDir::new()) } fn call_structured(structured: serde_json::Value) -> (ToolCallResult, Option) { @@ -366,7 +385,7 @@ mod tests { fn oversized_results_and_errors_keep_their_head_and_tail_and_save_the_full_text() { let text = (1..=100_000).map(|n| format!("{n}\n")).collect::>().concat(); let spill_dir = TempDir::new(); - let convert = |mcp| convert_tool_result(&req(), Ok(mcp), &spill_dir); + let convert = |mcp| convert_tool_result(&req(), Ok::<_, McpToolCallError>(mcp), &spill_dir); let (result, _) = convert(McpCallToolResult::success(vec![ContentBlock::text(&text)])).unwrap(); let error = convert(McpCallToolResult::error(vec![ContentBlock::text(&text)])).unwrap_err().error; diff --git a/crates/mcp-utils/src/client/tool_filter.rs b/crates/aether-core/src/mcp/tool_filter.rs similarity index 81% rename from crates/mcp-utils/src/client/tool_filter.rs rename to crates/aether-core/src/mcp/tool_filter.rs index 070ef3e5f..1ce23666c 100644 --- a/crates/mcp-utils/src/client/tool_filter.rs +++ b/crates/aether-core/src/mcp/tool_filter.rs @@ -1,4 +1,4 @@ -use llm::ToolDefinition; +use rmcp::model::Tool; use utils::matches_name_pattern; #[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] @@ -21,9 +21,9 @@ impl ToolMatcher { Self::Annotations(matcher) } - pub fn matches(&self, tool: &ToolDefinition) -> bool { + pub fn matches(&self, name: &str, tool: &Tool) -> bool { match self { - Self::Name(pattern) => matches_name_pattern(pattern, &tool.name), + Self::Name(pattern) => matches_name_pattern(pattern, name), Self::Annotations(matcher) => matcher.matches(tool), } } @@ -43,7 +43,7 @@ pub struct ToolAnnotationMatcher { } impl ToolAnnotationMatcher { - pub fn matches(&self, tool: &ToolDefinition) -> bool { + pub fn matches(&self, tool: &Tool) -> bool { let Some(annotations) = tool.annotations.as_ref() else { return false }; let pairs = [ (self.read_only, annotations.read_only_hint), @@ -58,7 +58,6 @@ impl ToolAnnotationMatcher { } } -/// Filter for restricting which MCP tools an agent may discover and execute. #[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct ToolFilter { @@ -73,13 +72,10 @@ impl ToolFilter { self.allow.is_empty() && self.deny.is_empty() } - pub fn apply(&self, tools: Vec) -> Vec { - tools.into_iter().filter(|tool| self.is_tool_allowed(tool)).collect() - } - - pub fn is_tool_allowed(&self, tool: &ToolDefinition) -> bool { - let allowed = self.allow.is_empty() || self.allow.iter().any(|matcher| matcher.matches(tool)); - let denied = self.deny.iter().any(|matcher| matcher.matches(tool)); + pub fn allows(&self, name: &str, tool: &Tool) -> bool { + let matches = |matcher: &ToolMatcher| matcher.matches(name, tool); + let allowed = self.allow.is_empty() || self.allow.iter().any(matches); + let denied = self.deny.iter().any(matches); allowed && !denied } } diff --git a/crates/aether-core/src/testing/fake_mcp.rs b/crates/aether-core/src/testing/fake_mcp.rs deleted file mode 100644 index 2b6ff2c2c..000000000 --- a/crates/aether-core/src/testing/fake_mcp.rs +++ /dev/null @@ -1,30 +0,0 @@ -use crate::mcp::{McpBuilder, ServerFactory}; -use futures::FutureExt; -use mcp_utils::client::{InMemoryServerSpec, McpServer, McpTransport, ToolExposure}; -use rmcp::{RoleServer, service::DynService}; - -pub use mcp_utils::testing::{ - CapturedTaskUpdate, CapturedToolCall, FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, fake_mcp, -}; - -pub trait McpBuilderTestExt { - fn with_fake_mcp(self, name: impl Into, server: FakeMcpServer) -> Self; -} - -impl McpBuilderTestExt for McpBuilder { - fn with_fake_mcp(self, name: impl Into, server: FakeMcpServer) -> Self { - let name = name.into(); - let factory_name = name.clone(); - let factory: ServerFactory = Box::new(move |_spec, _services| { - let server = server.clone(); - async move { Box::new(server) as Box> }.boxed() - }); - self.register_in_memory_server(factory_name.clone(), factory).with_servers(vec![McpServer::new( - name, - McpTransport::InMemory { - spec: InMemoryServerSpec { factory: factory_name, args: Vec::new(), input: None }, - }, - ToolExposure::ModelVisible, - )]) - } -} diff --git a/crates/aether-core/src/testing/mcp_test.rs b/crates/aether-core/src/testing/mcp_test.rs deleted file mode 100644 index 64adc7f90..000000000 --- a/crates/aether-core/src/testing/mcp_test.rs +++ /dev/null @@ -1,283 +0,0 @@ -use crate::events::{TaskOutcomeState, TraceContext, task_created_result}; -use crate::mcp::tool_bridge::{convert_tool_result, map_task_result_to_outcome}; -use crate::mcp::{McpHandle, McpRuntime, ServerFactory, ToolCallStream, mcp}; -use futures::{FutureExt, StreamExt}; -use mcp_utils::client::{ - CallToolOptions, CancellationToken, InMemoryServerSpec, McpConnectionDetails, McpServer, McpTransport, - ToolCallEvent, ToolExposure, ToolFilter, -}; -use mcp_utils::testing::{ElicitationScript, UrlElicitationHandler, url_elicitation_handler}; -use rmcp::model::{CreateTaskResult, ElicitResult, ProgressNotificationParam}; -use rmcp::{RoleServer, ServerHandler, service::DynService}; -use serde_json::Value; -use std::collections::{HashMap, VecDeque}; -use std::path::Path; -use std::sync::Mutex; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; -use tokio::sync::watch; -use utils::temp_dir::TempDir; - -pub use mcp_utils::testing::CapturedElicitation; - -const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(10); - -#[derive(Default)] -pub struct McpTestBuilder { - servers: Vec, - factories: Vec<(String, ServerFactory)>, - elicitation_responses: Vec, - on_url_elicitation: Option, - trace_context: Option, - tool_timeout: Duration, - tool_filter: ToolFilter, -} - -fn task_outcome(outcome: crate::events::TaskOutcome) -> TaskOutcome { - let (status, body) = match outcome.state { - TaskOutcomeState::Completed { result, .. } => ("completed", result.result), - TaskOutcomeState::Failed { error } => ("failed", error.error), - TaskOutcomeState::Cancelled => { - ("cancelled", "The background task was cancelled and will not produce a result.".into()) - } - }; - TaskOutcome { task_id: outcome.task_id, status: status.into(), body } -} - -pub struct McpTest { - mcp: McpHandle, - runtime: McpRuntime, - snapshot: McpConnectionDetails, - elicitations: ElicitationScript, - deferred_tools: tokio::sync::Mutex>, - cancel_tokens: Mutex>, - trace_context: Option, - tool_timeout: Duration, - next_call_id: AtomicU64, - spill_dir: TempDir, -} - -pub struct TaskOutcome { - pub task_id: String, - pub status: String, - pub body: String, -} - -pub struct ToolCallOutcome { - pub result: Result, - pub progress: Vec, - pub deferred_task: Option, -} - -struct DeferredTool { - request: llm::ToolCallRequest, - events: ToolCallStream, -} - -impl McpTestBuilder { - pub fn new() -> Self { - Self::default() - } - - pub fn server(self, name: impl Into, server: S) -> Self - where - S: ServerHandler + Clone + Send + Sync + 'static, - { - self.server_with_exposure(name, server, ToolExposure::ModelVisible) - } - - pub fn deferred_server(self, name: impl Into, server: T) -> Self - where - T: ServerHandler + Clone + Send + Sync + 'static, - { - self.server_with_exposure(name, server, ToolExposure::deferred_all()) - } - - pub fn server_with_exposure(mut self, name: impl Into, server: T, exposure: ToolExposure) -> Self - where - T: ServerHandler + Clone + Send + Sync + 'static, - { - let name = name.into(); - let factory_name = format!("test-{}", self.factories.len()); - let factory_server = server; - let factory: ServerFactory = Box::new(move |_spec, _services| { - let server = factory_server.clone(); - async move { Box::new(server) as Box> }.boxed() - }); - self.factories.push((factory_name.clone(), factory)); - self.servers.push(McpServer::new( - name, - McpTransport::InMemory { - spec: InMemoryServerSpec { factory: factory_name, args: Vec::new(), input: None }, - }, - exposure, - )); - self - } - - pub fn elicitation_response(mut self, response: ElicitResult) -> Self { - self.elicitation_responses.push(response); - self - } - - pub fn on_url_elicitation(mut self, handler: T) -> Self - where - T: Fn(String, String) -> Fut + Send + Sync + 'static, - Fut: Future + Send + 'static, - { - self.on_url_elicitation = Some(url_elicitation_handler(handler)); - self - } - - pub fn trace_context(mut self, trace_context: TraceContext) -> Self { - self.trace_context = Some(trace_context); - self - } - - pub fn tool_timeout(mut self, timeout: Duration) -> Self { - self.tool_timeout = timeout; - self - } - - /// Filter which tools the servers expose and the gateway serves. - pub fn tool_filter(mut self, filter: ToolFilter) -> Self { - self.tool_filter = filter; - self - } - - pub async fn build(self) -> McpTest { - let builder = mcp("/workspace").with_servers(self.servers).with_tool_filter(self.tool_filter); - let builder = self - .factories - .into_iter() - .fold(builder, |builder, (name, factory)| builder.register_in_memory_server(name, factory)); - let mut spawn = builder.spawn().await.expect("MCP test manager spawns"); - let snapshot = spawn.block_until_ready().await.expect("MCP test manager becomes ready"); - let (runtime, event_rx) = spawn.split(); - - McpTest { - mcp: runtime.handle().clone(), - runtime, - snapshot, - elicitations: ElicitationScript::spawn_with_url_handler( - event_rx, - self.elicitation_responses, - self.on_url_elicitation, - ), - deferred_tools: tokio::sync::Mutex::new(VecDeque::new()), - cancel_tokens: Mutex::new(HashMap::new()), - trace_context: self.trace_context, - tool_timeout: if self.tool_timeout.is_zero() { DEFAULT_TOOL_TIMEOUT } else { self.tool_timeout }, - next_call_id: AtomicU64::new(1), - spill_dir: TempDir::new(), - } - } -} - -impl McpTest { - pub async fn call(&self, server: &str, tool: &str, arguments: Value) -> ToolCallOutcome { - let id = self.next_call_id.fetch_add(1, Ordering::Relaxed); - let request = llm::ToolCallRequest { - id: format!("mcp-test-{id}"), - name: format!("{server}__{tool}"), - arguments: arguments.to_string(), - }; - let request_for_outcome = request.clone(); - let cancel = CancellationToken::new(); - self.cancel_tokens.lock().expect("cancel token lock").insert(request.id.clone(), cancel.clone()); - let options = CallToolOptions { - timeout: self.tool_timeout, - meta: self.trace_context.as_ref().map(TraceContext::to_meta), - cancel, - }; - let mut events = self.mcp.call_model_visible(request.name, &request.arguments, options); - - let mut progress = Vec::new(); - while let Some(event) = events.next().await { - match event { - ToolCallEvent::Progress(event) => progress.push(event), - ToolCallEvent::TaskCreated(task) => { - self.deferred_tools - .lock() - .await - .push_back(DeferredTool { request: request_for_outcome.clone(), events }); - return ToolCallOutcome { - result: Ok(task_created_result(&request_for_outcome, &task.task.task_id)), - progress, - deferred_task: Some(task), - }; - } - ToolCallEvent::Complete(outcome) => { - let result = - convert_tool_result(&request_for_outcome, outcome, &self.spill_dir).map(|(result, _)| result); - return ToolCallOutcome { result, progress, deferred_task: None }; - } - ToolCallEvent::TaskStatus(_) | ToolCallEvent::TaskComplete { .. } | ToolCallEvent::Cancelled { .. } => { - panic!("MCP task lifecycle event arrived before deferral") - } - } - } - panic!("MCP test tool event stream ended before completion"); - } - - pub fn cancel_tool(&self, tool_id: &str) { - let tokens = self.cancel_tokens.lock().expect("cancel token lock"); - tokens.get(tool_id).expect("cancel_tool targets a tool started with call()").cancel(); - } - - pub async fn next_tool_event(&self) -> Option { - self.next_deferred_event().await.map(|(_, event)| event) - } - - pub async fn next_task_outcome(&self) -> Option { - while let Some((request, event)) = self.next_deferred_event().await { - match event { - ToolCallEvent::TaskComplete { task, result } => { - return Some(task_outcome(map_task_result_to_outcome(request, task, result, &self.spill_dir))); - } - ToolCallEvent::Cancelled { task_id } => { - return Some(task_outcome(crate::events::TaskOutcome { - request, - task_id: task_id.unwrap_or_else(|| "pending".to_string()), - state: TaskOutcomeState::Cancelled, - })); - } - ToolCallEvent::Progress(_) - | ToolCallEvent::TaskCreated(_) - | ToolCallEvent::TaskStatus(_) - | ToolCallEvent::Complete(_) => {} - } - } - None - } - - async fn next_deferred_event(&self) -> Option<(llm::ToolCallRequest, ToolCallEvent)> { - let mut deferred = self.deferred_tools.lock().await; - loop { - let front = deferred.front_mut()?; - match front.events.next().await { - Some(event) => return Some((front.request.clone(), event)), - None => { - deferred.pop_front(); - } - } - } - } - - pub fn snapshot(&self) -> &McpConnectionDetails { - &self.snapshot - } - - /// The Unix socket path of the deferred-tool gateway, if one was started. - pub fn gateway_endpoint(&self) -> Option<&Path> { - self.runtime.gateway_endpoint() - } - - pub fn subscribe(&self) -> watch::Receiver { - self.mcp.subscribe() - } - - pub fn elicitations(&self) -> Vec { - self.elicitations.captured() - } -} diff --git a/crates/aether-core/src/testing/mod.rs b/crates/aether-core/src/testing/mod.rs index 1816d4ae3..de021f734 100644 --- a/crates/aether-core/src/testing/mod.rs +++ b/crates/aether-core/src/testing/mod.rs @@ -1,13 +1,9 @@ mod agent_event_builder; mod agent_trace; mod fake_agent_observer; -mod fake_mcp; -mod mcp_test; mod utils; pub use agent_event_builder::*; pub use agent_trace::*; pub use fake_agent_observer::*; -pub use fake_mcp::*; -pub use mcp_test::*; pub use utils::*; diff --git a/crates/aether-core/src/testing/utils.rs b/crates/aether-core/src/testing/utils.rs index cccf4ef41..fa3616886 100644 --- a/crates/aether-core/src/testing/utils.rs +++ b/crates/aether-core/src/testing/utils.rs @@ -1,15 +1,20 @@ -use std::collections::BTreeMap; +use mcp_utils::client::ElicitationSink; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::ElicitResult; +use mcp_utils::server::McpServer; +use mcp_utils::testing::{CapturedElicitation, ElicitationScriptBuilder, FakeMcpServer}; +use std::future::Future; use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::sync::{Notify, mpsc}; use crate::context::CompactionConfig; -use crate::core::{AgentError, Prompt, RetryConfig, agent}; +use crate::core::{AgentError, McpInstructions, Prompt, RetryConfig, agent}; use crate::events::{ AgentCommand, AgentEvent, AgentObserver, Command, ContextEvent, ToolEvent, TurnEvent, UserCommand, }; -use crate::mcp::mcp; -use crate::testing::{AgentTrace, FakeAgentObserver, FakeMcpServer, McpBuilderTestExt}; +use crate::mcp::{McpBuilder, McpCatalog, mcp}; +use crate::testing::{AgentTrace, FakeAgentObserver}; use llm::{ChatMessage, Context, LlmError, LlmModel, LlmResponse, ModelSettings, StreamingModelProvider}; use llm::testing::FakeLlmProvider; @@ -47,8 +52,15 @@ pub fn content_events(events: Vec) -> Vec { .collect() } -pub fn mcp_instructions(entries: &[(&str, &str)]) -> BTreeMap { - entries.iter().map(|(k, v)| ((*k).to_string(), (*v).to_string())).collect() +pub async fn ready_catalog(builder: McpBuilder) -> McpCatalog { + let catalog = builder.spawn().expect("the MCP catalog spawns"); + catalog.ready().await; + catalog +} + +pub fn mcp_instructions(entries: &[(&str, &str)]) -> McpInstructions { + let servers = entries.iter().map(|(k, v)| ((*k).to_string(), (*v).to_string())).collect(); + McpInstructions { servers, progressive_discovery: false } } /// Millisecond-scale retry delays so retry tests run fast under virtual time. @@ -182,10 +194,11 @@ impl From> for TestScenario { } } -/// Result of running a test agent, including messages and captured contexts. +/// Result of running a test agent, including messages, captured contexts, and the MCP elicitations it answered. pub struct TestAgentResult { pub messages: Vec, pub captured_contexts: Arc>>, + pub elicitations: Vec, } /// Error returned by [`TestAgentBuilder::run`], [`TestAgentBuilder::run_trace`], and @@ -213,7 +226,9 @@ struct AgentTestConfig { max_auto_continues: Option, retry_config: Option, observers: Vec>, - mcp_server: Option<(String, FakeMcpServer)>, + /// `None` connects a single fake server named `test`. + mcp_servers: Option>, + elicitations: Option, initial_messages: Vec, system_prompts: Vec, session_affinity_key: Option, @@ -248,7 +263,8 @@ impl TestAgentBuilder { max_auto_continues: None, retry_config: None, observers: Vec::new(), - mcp_server: Some(("test".to_string(), FakeMcpServer::new())), + mcp_servers: None, + elicitations: None, initial_messages: Vec::new(), system_prompts: Vec::new(), session_affinity_key: None, @@ -313,13 +329,26 @@ impl TestAgentBuilder { /// Run without the default fake MCP server when the scenario does not exercise tools. pub fn without_mcp(mut self) -> Self { - self.agent.mcp_server = None; + self.agent.mcp_servers = Some(Vec::new()); + self + } + + pub fn mcp_server(mut self, name: &str, server: impl Into) -> Self { + self.agent.mcp_servers.get_or_insert_with(Vec::new).push(ServerSpec::in_process(name, server)); + self + } + + pub fn elicitation_response(mut self, response: ElicitResult) -> Self { + self.agent.elicitations = Some(self.agent.elicitations.take().unwrap_or_default().response(response)); self } - /// Replace the default fake MCP server with a scripted server. - pub fn fake_mcp_server(mut self, name: &str, server: FakeMcpServer) -> Self { - self.agent.mcp_server = Some((name.to_string(), server)); + pub fn on_url_elicitation(mut self, handler: F) -> Self + where + F: Fn(String, String) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + self.agent.elicitations = Some(self.agent.elicitations.take().unwrap_or_default().on_url(handler)); self } @@ -397,18 +426,17 @@ impl TestAgentBuilder { } let captured_contexts = llm.captured_contexts(); - let mut mcp_spawn = match config.mcp_server { - Some((name, server)) => { - Some(mcp("/workspace").with_fake_mcp(name, server).spawn().await.map_err(AgentError::from)?) - } - None => None, - }; - - let mut builder = agent(llm); - if let Some(spawn) = &mut mcp_spawn { - let snapshot = spawn.block_until_ready().await.expect("bootstrap completes"); - builder = builder.tools(spawn.handle().clone(), snapshot.tool_definitions()); + let servers = config.mcp_servers.unwrap_or_else(|| vec![ServerSpec::in_process("test", FakeMcpServer::math())]); + let mut mcp = mcp("/workspace").with_servers(servers); + let mut elicitations = None; + if let Some(script) = config.elicitations { + let (sink, events) = mpsc::channel(32); + mcp = mcp.with_elicitation(ElicitationSink::form_and_url(sink)); + elicitations = Some(script.spawn(events)); } + let mcp = ready_catalog(mcp).await; + + let mut builder = agent(llm).mcp(mcp); if let Some(timeout) = config.timeout { builder = builder.tool_timeout(timeout); } @@ -473,7 +501,8 @@ impl TestAgentBuilder { handle.await_completion().await; - Ok(TestAgentResult { messages, captured_contexts }) + let elicitations = elicitations.map(|script| script.captured()).unwrap_or_default(); + Ok(TestAgentResult { messages, captured_contexts, elicitations }) } fn with_execution(mut self, execution: TestExecution) -> Self { diff --git a/crates/aether-core/tests/agent/agent_tests.rs b/crates/aether-core/tests/agent/agent_tests.rs index 6604a1080..2f2a26031 100644 --- a/crates/aether-core/tests/agent/agent_tests.rs +++ b/crates/aether-core/tests/agent/agent_tests.rs @@ -8,6 +8,7 @@ use aether_core::{ }; use llm::{ChatMessage, ContentBlock, StopReason}; use llm::{MessageId, testing::llm_response}; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse}; use serde_json::json; fn completed_message_ids(events: &[AgentEvent]) -> Vec { @@ -117,13 +118,36 @@ async fn failed_mcp_command_submission_completes_the_tool_with_an_error() -> Res matches!( event, AgentEvent::Tool(ToolEvent::Error { error }) - if error.id == "call_1" && error.error.contains("Failed to resolve tool") + if error.id == "call_1" && error.error.contains("Tool not found: test__missing_tool") ) })); assert!(matches!(messages.last(), Some(AgentEvent::Turn(TurnEvent::Ended { .. })))); Ok(()) } +#[tokio::test] +async fn malformed_tool_arguments_fail_the_tool_without_reaching_the_server() -> Result<(), Box> { + let server = FakeMcpServer::math(); + let state = server.state(); + let llm_responses = [ + llm_response().tool_call("call_1", "test__add_numbers", &["{not json"]).build(), + llm_response().text(&["recovered"]).build(), + ]; + + let messages = test_agent().mcp_server("test", server).llm_responses(&llm_responses).user_text("add").run().await?; + + assert!(messages.iter().any(|event| { + matches!( + event, + AgentEvent::Tool(ToolEvent::Error { error }) + if error.id == "call_1" && error.error.starts_with("Invalid tool arguments") + ) + })); + assert!(state.calls_for("add_numbers").is_empty()); + assert!(matches!(messages.last(), Some(AgentEvent::Turn(TurnEvent::Ended { .. })))); + Ok(()) +} + #[tokio::test] async fn streamed_tool_input_is_followed_by_the_complete_call() -> Result<(), Box> { let tool_request = json!({ "a": 3, "b": 5 }); @@ -246,21 +270,14 @@ async fn test_cancellation() -> Result<(), Box> { #[tokio::test(start_paused = true)] async fn test_tool_timeout() -> Result<(), Box> { - let tool_duration = 2000; - let tool_timeout = 500; - - let tool_request = json!({ "sleep_ms": tool_duration }); - let (t1_id, t1_name) = ("call_1", "test__slow_tool"); - - let llm_responses = [ - llm_response().tool_call(t1_id, t1_name, &[&tool_request.to_string()]).build(), - llm_response().text(&["done"]).build(), - ]; + let llm_responses = + [llm_response().tool_call("call_1", "test__slow", &["{}"]).build(), llm_response().text(&["done"]).build()]; let messages = test_agent() + .mcp_server("test", FakeMcpServer::new().with_tool(FakeTool::new("slow").responds(FakeToolResponse::pending()))) .llm_responses(&llm_responses) .user_text("run slow tool") - .tool_timeout(Duration::from_millis(tool_timeout)) + .tool_timeout(Duration::from_millis(500)) .run() .await?; diff --git a/crates/aether-core/tests/agent/cancel_tests.rs b/crates/aether-core/tests/agent/cancel_tests.rs index a871387ac..39d40ba0a 100644 --- a/crates/aether-core/tests/agent/cancel_tests.rs +++ b/crates/aether-core/tests/agent/cancel_tests.rs @@ -1,19 +1,18 @@ use aether_core::events::{AgentEvent, MessageEvent, ToolEvent, TurnEvent}; -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, TestScenario, test_agent}; +use aether_core::testing::{TestScenario, test_agent}; use llm::testing::llm_response; -use rmcp::model::{CallToolResult, ContentBlock, CreateTaskResult, DetailedTask, Task, TaskPayload, TaskStatus}; +use mcp_utils::testing::{FakeMcpServer, FakeTask, FakeTool, FakeToolResponse, completed_task, task, working_task}; +use rmcp::model::TaskStatus; use std::sync::Arc; use tokio::sync::Notify; #[tokio::test] async fn user_cancel_does_not_report_foreground_tool_as_background_task() { - let server = FakeMcpServer::new().with_tool( - FakeTool::new("slow").responds(FakeToolResponse::text("late result").delay(std::time::Duration::from_secs(1))), - ); + let server = FakeMcpServer::new().with_tool(FakeTool::new("slow").responds(FakeToolResponse::pending())); let arguments = serde_json::json!({}).to_string(); let events = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("foreground-call", "tasks__slow", &[&arguments]).build(), llm_response().text(&["still works"]).build(), @@ -39,25 +38,14 @@ async fn user_cancel_does_not_report_foreground_tool_as_background_task() { #[tokio::test] async fn immediate_background_outcome_follows_deferred_event() { - let now = chrono::Utc::now().to_rfc3339(); - let seed = Task::new("terminal-task", TaskStatus::Completed, now.clone(), now.clone()); - let completed = Task::new("terminal-task", TaskStatus::Completed, now.clone(), now); - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text("finished")]); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("terminal").responds(FakeToolResponse::task(CreateTaskResult::new(seed)))) - .with_task( - "terminal-task", - [DetailedTask::new( - completed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap(), - }, - )], - ); + let server = FakeTask::new("terminal-task") + .seed(task("terminal-task", TaskStatus::Completed)) + .states([completed_task("terminal-task", "finished")]) + .server("terminal"); let arguments = serde_json::json!({}).to_string(); let events = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("terminal-call", "tasks__terminal", &[&arguments]).build(), llm_response().text(&["background task started"]).build(), @@ -87,32 +75,14 @@ async fn immediate_background_outcome_follows_deferred_event() { #[tokio::test] async fn background_task_progress_reaches_agent_events() { - let now = chrono::Utc::now().to_rfc3339(); - let seed = Task::new("progress-task", TaskStatus::Working, now.clone(), now.clone()).with_poll_interval_ms(10); - let working = Task::new("progress-task", TaskStatus::Working, now.clone(), now.clone()).with_poll_interval_ms(10); - let completed = Task::new("progress-task", TaskStatus::Completed, now.clone(), now); - let result = CallToolResult::success(vec![ContentBlock::text("finished")]); - let server = FakeMcpServer::new() - .with_tool( - FakeTool::new("deferred") - .responds(FakeToolResponse::task(CreateTaskResult::new(seed)).task_progress(1.0, Some(2.0))), - ) - .with_task( - "progress-task", - [ - DetailedTask::new(working, TaskPayload::Working), - DetailedTask::new( - completed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap(), - }, - ), - ], - ); + let server = FakeTask::new("progress-task") + .progress(1.0, Some(2.0)) + .states([working_task("progress-task"), completed_task("progress-task", "finished")]) + .server("deferred"); let arguments = serde_json::json!({}).to_string(); let events = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("progress-call", "tasks__deferred", &[&arguments]).build(), llm_response().text(&["background task started"]).build(), @@ -142,16 +112,12 @@ async fn background_task_progress_reaches_agent_events() { #[tokio::test] async fn closing_agent_input_cancels_background_tasks() { - let now = chrono::Utc::now().to_rfc3339(); - let task = Task::new("shutdown-task", TaskStatus::Working, now.clone(), now).with_poll_interval_ms(10); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(task.clone())))) - .with_task("shutdown-task", [DetailedTask::new(task, TaskPayload::Working)]); + let server = FakeTask::new("shutdown-task").server("deferred"); let server_state = server.state(); let arguments = serde_json::json!({}).to_string(); test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("shutdown-call", "tasks__deferred", &[&arguments]).build(), llm_response().text(&["background task started"]).build(), @@ -171,17 +137,13 @@ async fn closing_agent_input_cancels_background_tasks() { #[tokio::test] async fn user_cancel_surfaces_background_task_cancellation() { - let now = chrono::Utc::now().to_rfc3339(); - let task = Task::new("cancel-task", TaskStatus::Working, now.clone(), now).with_poll_interval_ms(10); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(task.clone())))) - .with_task("cancel-task", [DetailedTask::new(task, TaskPayload::Working)]); + let server = FakeTask::new("cancel-task").server("deferred"); let server_state = server.state(); let arguments = serde_json::json!({}).to_string(); let release = Arc::new(Notify::new()); let events = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("cancel-call", "tasks__deferred", &[&arguments]).build(), llm_response().build(), @@ -207,27 +169,12 @@ async fn user_cancel_surfaces_background_task_cancellation() { #[tokio::test] async fn user_cancel_preserves_queued_background_task_outcome() { - let now = chrono::Utc::now().to_rfc3339(); - let initial_task = - Task::new("queued-task", TaskStatus::Working, now.clone(), now.clone()).with_poll_interval_ms(10); - let completed = Task::new("queued-task", TaskStatus::Completed, now.clone(), now); - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text("finished")]); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(initial_task)))) - .with_task( - "queued-task", - [DetailedTask::new( - completed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap(), - }, - )], - ); + let server = FakeTask::new("queued-task").states([completed_task("queued-task", "finished")]).server("deferred"); let arguments = serde_json::json!({}).to_string(); let release = Arc::new(Notify::new()); let events = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("queued-call", "tasks__deferred", &[&arguments]).build(), llm_response().build(), diff --git a/crates/aether-core/tests/agent/clear_context_tests.rs b/crates/aether-core/tests/agent/clear_context_tests.rs index d87ab793d..a353f6c81 100644 --- a/crates/aether-core/tests/agent/clear_context_tests.rs +++ b/crates/aether-core/tests/agent/clear_context_tests.rs @@ -1,25 +1,21 @@ use aether_core::core::Prompt; use aether_core::events::{AgentEvent, Command, ContextEvent, ToolEvent, UserCommand}; -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, TestScenario, test_agent}; +use aether_core::testing::{TestScenario, test_agent}; use llm::testing::llm_response; use llm::{ChatMessage, ContentBlock}; -use rmcp::model::{CreateTaskResult, DetailedTask, Task, TaskPayload, TaskStatus}; +use mcp_utils::testing::FakeTask; use std::sync::Arc; use tokio::sync::Notify; #[tokio::test] async fn clear_context_suppresses_cancelled_background_task_notification() -> Result<(), Box> { - let now = chrono::Utc::now().to_rfc3339(); - let task = Task::new("clear-task", TaskStatus::Working, now.clone(), now).with_poll_interval_ms(10); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(task.clone())))) - .with_task("clear-task", [DetailedTask::new(task, TaskPayload::Working)]); + let server = FakeTask::new("clear-task").server("deferred"); let server_state = server.state(); let arguments = serde_json::json!({}).to_string(); let release = Arc::new(Notify::new()); let result = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .llm_responses(&[ llm_response().tool_call("clear-call", "tasks__deferred", &[&arguments]).build(), llm_response().text(&["must not finish"]).build(), diff --git a/crates/aether-core/tests/agent/prompt_cache_tests.rs b/crates/aether-core/tests/agent/prompt_cache_tests.rs index f2255654d..d192e7d8c 100644 --- a/crates/aether-core/tests/agent/prompt_cache_tests.rs +++ b/crates/aether-core/tests/agent/prompt_cache_tests.rs @@ -1,10 +1,11 @@ +use mcp_utils::testing::{FakeMcpServer, FakeTool}; use std::error::Error; use aether_core::core::Prompt; -use aether_core::events::{AgentCommand, Command}; -use aether_core::testing::{FakeMcpServer, FakeTool, TestScenario, test_agent}; +use aether_core::events::{AgentEvent, ToolEvent}; +use aether_core::testing::{TestScenario, test_agent}; +use llm::LlmResponse; use llm::testing::llm_response; -use llm::{LlmResponse, ToolDefinition}; #[tokio::test] async fn derives_same_cache_key_for_shared_prompt_prefix() -> Result<(), Box> { @@ -35,15 +36,20 @@ async fn cache_key_changes_with_tools() -> Result<(), Box> { #[tokio::test] async fn cache_key_refreshes_after_tool_updates() -> Result<(), Box> { + let server = FakeMcpServer::new(); + let state = server.state(); let result = test_agent() - .without_mcp() + .mcp_server("test", server) .system_prompt(Prompt::text("system prompt")) .llm_responses(&[response(), response()]) .scenario( TestScenario::new() .user_text("first question") .wait_for_turn_end() - .send(Command::AgentCommand(AgentCommand::UpdateTools(vec![tool("read_file")]))) + .perform(move || { + tokio::spawn(async move { state.add_tool_and_notify(FakeTool::new("read_file")).await }); + }) + .wait_for(|event| defines_tool(event, "test__read_file")) .user_text("second question") .wait_for_turn_end(), ) @@ -126,7 +132,7 @@ async fn prompt_cache_key( user_prompt: &str, ) -> Result> { let builder = match tools { - Some(server) => test_agent().fake_mcp_server("test", server), + Some(server) => test_agent().mcp_server("test", server), None => test_agent().without_mcp(), }; let result = builder @@ -170,6 +176,6 @@ fn server_with_tool(name: &str) -> FakeMcpServer { FakeMcpServer::new().with_tool(FakeTool::new(name).description(format!("{name} description"))) } -fn tool(name: &str) -> ToolDefinition { - ToolDefinition::new(name, format!("{name} description"), serde_json::json!({ "type": "object" })) +fn defines_tool(event: &AgentEvent, name: &str) -> bool { + matches!(event, AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools }) if tools.iter().any(|tool| tool.name == name)) } diff --git a/crates/aether-core/tests/agent/retry_tests.rs b/crates/aether-core/tests/agent/retry_tests.rs index 1345aa267..896a6b6c4 100644 --- a/crates/aether-core/tests/agent/retry_tests.rs +++ b/crates/aether-core/tests/agent/retry_tests.rs @@ -1,13 +1,13 @@ use aether_core::events::TurnEvent; +use mcp_utils::testing::{FakeTask, Gate}; use std::error::Error; use std::time::Duration; use aether_core::core::RetryConfig; use aether_core::events::{AgentEvent, LlmCallOutcome, TurnOutcome}; -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, fast_retry, test_agent}; +use aether_core::testing::{TestScenario, fast_retry, test_agent}; use llm::ProviderError; use llm::testing::{failed_call, llm_response}; -use rmcp::model::{CreateTaskResult, DetailedTask, Task, TaskPayload, TaskStatus}; fn retry_attempts(messages: &[AgentEvent]) -> Vec { messages @@ -33,21 +33,17 @@ async fn deferred_event_after_retry_clears_pending_tool_and_cancels_task() -> Re llm_response().text(&["recovered"]).build_results(), ]; - let now = chrono::Utc::now().to_rfc3339(); - let task = Task::new("stale-task", TaskStatus::Working, now.clone(), now).with_poll_interval_ms(10); - let server = FakeMcpServer::new() - .with_tool( - FakeTool::new("deferred") - .responds(FakeToolResponse::task(CreateTaskResult::new(task.clone())).delay(Duration::from_millis(10))), - ) - .with_task("stale-task", [DetailedTask::new(task, TaskPayload::Working)]); + let created = Gate::new(); + let server = FakeTask::new("stale-task").gated(created.clone()).server("deferred"); let server_state = server.state(); let result = test_agent() - .fake_mcp_server("tasks", server) + .mcp_server("tasks", server) .retry_config(fast_retry(1)) .llm_result_responses(&attempts) - .user_text("go") + .scenario( + TestScenario::new().user_text("go").wait_for_retry(1).perform(move || created.open()).wait_for_turn_end(), + ) .run_with_context() .await?; diff --git a/crates/aether-core/tests/agent/usage_tests.rs b/crates/aether-core/tests/agent/usage_tests.rs index 2b99b741a..f3a95f60d 100644 --- a/crates/aether-core/tests/agent/usage_tests.rs +++ b/crates/aether-core/tests/agent/usage_tests.rs @@ -1,9 +1,10 @@ use aether_core::context::CompactionConfig; use aether_core::events::{AgentEvent, LlmCallOutcome, SubAgentProgressPayload, ToolEvent, TurnEvent, TurnOutcome}; -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, TestScenario, fast_retry, test_agent}; +use aether_core::testing::{TestScenario, fast_retry, test_agent}; use llm::alloyed::AlloyedModelProvider; use llm::testing::{FakeLlmProvider, llm_response, priced_model, session_usage_event}; use llm::{ChatMessage, LlmCallPurpose, LlmModel, ProviderError, SessionUsageEvent, TokenUsage, UsageSource, Usd}; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse}; use std::sync::Arc; use tokio::sync::Notify; @@ -181,7 +182,7 @@ async fn sub_agent_usage_is_folded_into_the_parent_totals_before_the_tool_result let arguments = serde_json::json!({}).to_string(); let events = test_agent() - .fake_mcp_server("agents", server) + .mcp_server("agents", server) .llm_responses(&[ llm_response().tool_call("spawn-call", "agents__spawn", &[&arguments]).build(), llm_response().usage(1, 1).text(&["all done"]).build(), diff --git a/crates/aether-core/tests/doc_examples.rs b/crates/aether-core/tests/doc_examples.rs index a7c21a021..56a1b85d1 100644 --- a/crates/aether-core/tests/doc_examples.rs +++ b/crates/aether-core/tests/doc_examples.rs @@ -2,8 +2,8 @@ //! These `.md` files are rendered into the auto-generated settings reference. use aether_core::core::PromptSource; +use aether_core::mcp::ToolFilter; use aether_doctest::assert_json_examples; -use mcp_utils::client::ToolFilter; #[test] fn prompt_source_examples_parse() { diff --git a/crates/aether-core/tests/mcp.rs b/crates/aether-core/tests/mcp.rs index 38bf2349c..7f14de492 100644 --- a/crates/aether-core/tests/mcp.rs +++ b/crates/aether-core/tests/mcp.rs @@ -4,8 +4,10 @@ mod mcp { mod instructions_tests; mod ipc_tests; mod mrtr_tests; - mod oauth_tests; mod server_round_trip_tests; + mod support; mod task_tests; + mod tool_catalog_tests; + mod tool_filter_tests; mod trace_context_tests; } diff --git a/crates/aether-core/tests/mcp/agent_sync_tests.rs b/crates/aether-core/tests/mcp/agent_sync_tests.rs index 1624406a3..3513c0ca0 100644 --- a/crates/aether-core/tests/mcp/agent_sync_tests.rs +++ b/crates/aether-core/tests/mcp/agent_sync_tests.rs @@ -1,57 +1,44 @@ use aether_core::{ - core::AgentDeps, - events::{AgentCommand, Command}, - mcp::{ServerFactory, mcp}, - testing::{FakeMcpServer, FakeTool}, + core::agent, + events::{AgentEvent, ToolEvent}, + mcp::mcp, + testing::{TestScenario, test_agent}, }; -use futures::FutureExt; -use mcp_utils::client::{InMemoryServerSpec, McpClientEvent, McpServer, McpTransport}; -use rmcp::RoleServer; -use rmcp::model::{ClientCapabilities, UrlElicitationCapability}; -use rmcp::service::DynService; +use llm::testing::{FakeLlmProvider, llm_response}; +use llm::{ChatMessage, Context, LlmResponse}; +use mcp_utils::client::ElicitationSink; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::{ElicitationCapability, UrlElicitationCapability}; +use mcp_utils::testing::{FakeMcpServer, FakeTool}; +use std::error::Error; use tokio::sync::mpsc; #[tokio::test] -async fn session_synchronization_does_not_keep_agent_input_open() { - let session = mcp("/workspace").spawn().await.unwrap(); - let (agent_tx, agent_rx) = mpsc::channel(32); - let (runtime, _) = session.connect_agent(agent_tx.clone()).await.split(); - - assert_eq!(agent_rx.sender_strong_count(), 1); - assert_eq!(agent_rx.sender_weak_count(), 1); +async fn an_agent_following_a_gateway_still_stops_once_its_input_closes() { + let mcp = + mcp("/workspace").with_servers(vec![ServerSpec::in_process("math", FakeMcpServer::new())]).spawn().unwrap(); + let (agent_tx, _agent_rx, handle) = agent(FakeLlmProvider::new(Vec::new())).mcp(mcp).spawn().await.unwrap(); drop(agent_tx); - assert_eq!(agent_rx.sender_strong_count(), 0); - drop(runtime); + + handle.await_completion().await; } #[tokio::test] async fn spawned_mcp_client_advertises_the_configured_elicitation_support() { let server = FakeMcpServer::new(); let state = server.state(); - let factory_server = server.clone(); - let factory: ServerFactory = Box::new(move |_, _| { - let server = factory_server.clone(); - async move { Box::new(server) as Box> }.boxed() - }); - let configured = McpServer::new( - "capability-capture", - McpTransport::InMemory { - spec: InMemoryServerSpec { factory: "capability-factory".to_string(), args: Vec::new(), input: None }, - }, - mcp_utils::client::ToolExposure::ModelVisible, - ); - let mut url_only = ClientCapabilities::builder().enable_elicitation().build(); - url_only.elicitation.as_mut().unwrap().url = Some(UrlElicitationCapability::default()); - let mut session = mcp("/workspace") - .register_in_memory_server("capability-factory", factory) - .with_servers(vec![configured]) - .with_agent_deps(AgentDeps::default().with_mcp_client_capabilities(url_only)) + let (sink, _host) = mpsc::channel(1); + let mcp = mcp("/workspace") + .with_servers(vec![ServerSpec::in_process("capability-capture", server)]) + .with_elicitation(ElicitationSink::new( + sink, + ElicitationCapability::new().with_url(UrlElicitationCapability::new()), + )) .spawn() - .await .unwrap(); - session.block_until_ready().await.expect("MCP bootstrap completes"); + mcp.ready().await; let capabilities = state.client_capabilities().expect("client capabilities were discovered"); let elicitation = capabilities.elicitation.expect("elicitation is advertised"); @@ -60,84 +47,57 @@ async fn spawned_mcp_client_advertises_the_configured_elicitation_support() { } #[tokio::test] -async fn session_synchronizes_agent_while_forwarding_host_events() { - let server = FakeMcpServer::new(); - let state = server.state(); - let factory_server = server.clone(); - let factory: ServerFactory = Box::new(move |_, _| { - let server = factory_server.clone(); - async move { Box::new(server) as Box> }.boxed() - }); - let configured = McpServer::new( - "dynamic", - McpTransport::InMemory { - spec: InMemoryServerSpec { factory: "dynamic-factory".to_string(), args: Vec::new(), input: None }, - }, - mcp_utils::client::ToolExposure::ModelVisible, - ); - let session = mcp("/workspace") - .register_in_memory_server("dynamic-factory", factory) - .with_servers(vec![configured]) - .spawn() - .await - .unwrap(); - let handle = session.handle().clone(); - let (agent_tx, mut agent_rx) = mpsc::channel(32); - let (runtime, mut host_events) = session.connect_agent(agent_tx.clone()).await.split(); +async fn the_agent_follows_the_gateway_catalog_as_it_changes() -> Result<(), Box> { + let server = FakeMcpServer::math().with_instructions("Use the math tools"); + let (adding, clearing) = (server.state(), server.state()); + let result = test_agent() + .mcp_server("dynamic", server) + .llm_responses(&[response(), response(), response()]) + .scenario( + TestScenario::new() + .user_text("initially") + .wait_for_turn_end() + .perform(move || { + tokio::spawn(async move { adding.add_tool_and_notify(FakeTool::new("added_later")).await }); + }) + .wait_for(|event| defines(event, |tools| tools.contains(&"dynamic__added_later"))) + .user_text("after adding a tool") + .wait_for_turn_end() + .perform(move || { + tokio::spawn(async move { clearing.clear_tools_and_notify().await }); + }) + .wait_for(|event| defines(event, |tools| tools.is_empty())) + .user_text("after clearing the tools") + .wait_for_turn_end(), + ) + .run_with_context() + .await?; - let mut received_initial_tools = false; - let mut received_initial_instructions = false; - while !received_initial_tools || !received_initial_instructions { - match agent_rx.recv().await.expect("agent synchronization remains connected") { - Command::AgentCommand(AgentCommand::UpdateTools(tools)) => { - received_initial_tools = tools.iter().any(|tool| tool.name == "dynamic__add_numbers"); - } - Command::AgentCommand(AgentCommand::UpdateMcpInstructions { server, body }) if server == "dynamic" => { - received_initial_instructions = body.as_deref() == Some("A fake MCP server for testing"); - } - _ => {} - } - } + let contexts = result.captured_contexts.lock().unwrap(); + assert!(tool_names(&contexts[0]).contains(&"dynamic__add_numbers".to_string())); + assert!(system_prompt(&contexts[0]).contains("Use the math tools")); + assert!(tool_names(&contexts[1]).contains(&"dynamic__added_later".to_string())); + assert_eq!(tool_names(&contexts[2]).len(), 0); + assert!(!system_prompt(&contexts[2]).contains("Use the math tools")); + Ok(()) +} - let mut received_status = false; - loop { - match host_events.recv().await.expect("host event stream remains connected") { - McpClientEvent::ServerStatusesChanged(_) => received_status = true, - McpClientEvent::ConnectionReady(_) => break, - McpClientEvent::Elicitation(_) - | McpClientEvent::ElicitationComplete { .. } - | McpClientEvent::AuthenticationFailed { .. } => {} - } - } - assert!(received_status); +fn defines(event: &AgentEvent, check: impl Fn(&[&str]) -> bool) -> bool { + let AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools }) = event else { return false }; + check(&tools.iter().map(|tool| tool.name.as_str()).collect::>()) +} - let mut snapshots = handle.subscribe(); - snapshots.borrow_and_update(); - state.add_tool_and_notify(FakeTool::new("added_later")).await; - snapshots.changed().await.expect("tool refresh publishes a snapshot"); - loop { - if let Command::AgentCommand(AgentCommand::UpdateTools(tools)) = - agent_rx.recv().await.expect("agent synchronization remains connected") - && tools.iter().any(|tool| tool.name == "dynamic__added_later") - { - break; - } - } +fn tool_names(context: &Context) -> Vec { + context.tools().iter().map(|tool| tool.name.clone()).collect() +} - snapshots.borrow_and_update(); - state.clear_tools_and_notify().await; - snapshots.changed().await.expect("tool removal publishes a snapshot"); - let mut received_empty_tools = false; - let mut received_removed_instructions = false; - while !received_empty_tools || !received_removed_instructions { - match agent_rx.recv().await.expect("agent synchronization remains connected") { - Command::AgentCommand(AgentCommand::UpdateTools(tools)) => received_empty_tools = tools.is_empty(), - Command::AgentCommand(AgentCommand::UpdateMcpInstructions { server, body }) if server == "dynamic" => { - received_removed_instructions = body.is_none(); - } - _ => {} - } +fn system_prompt(context: &Context) -> String { + match context.messages().first() { + Some(ChatMessage::System { content, .. }) => content.clone(), + _ => String::new(), } +} - drop(runtime); +fn response() -> Vec { + llm_response().text(&["done"]).build() } diff --git a/crates/aether-core/tests/mcp/config_parser_tests.rs b/crates/aether-core/tests/mcp/config_parser_tests.rs index ba71cf084..8616dbea5 100644 --- a/crates/aether-core/tests/mcp/config_parser_tests.rs +++ b/crates/aether-core/tests/mcp/config_parser_tests.rs @@ -1,201 +1,380 @@ -use mcp_utils::client::{McpConfig, McpHttpConfig, McpServer, McpTransport, ParseError}; -use reqwest::header::AUTHORIZATION; -use std::env; +use aether_auth::OAuthClientRegistration; +use aether_core::mcp::{ + McpConfig, McpOAuthConfig, McpServerConfig, ParseError, StdioServerConfig, ToolExposure, ToolFilter, ToolMatcher, + http_transport, +}; +use http::header::{AUTHORIZATION, HeaderMap, HeaderValue}; +use mcp_utils::client::{OAuthSettings, Transport}; +use mcp_utils::model::{Tool, ToolAnnotations}; use std::num::NonZeroU16; -use utils::variables::Vars; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tempfile::tempdir; +use utils::variables::{VarError, Vars}; -fn parse_servers(json: &str) -> Result, ParseError> { - McpConfig::from_json(json).unwrap().into_servers(&Vars::new()) +#[test] +fn from_json_accepts_mcp_servers_key() { + let config = McpConfig::from_json(r#"{"mcpServers": {"alpha": {"type": "stdio", "command": "a"}}}"#).unwrap(); + + assert_eq!(config.servers.keys().collect::>(), ["alpha"]); } -fn parse_one(json: &str) -> McpServer { - let mut servers = parse_servers(json).unwrap(); - assert_eq!(servers.len(), 1); - servers.remove(0) +#[test] +fn from_json_defaults_missing_type_to_stdio() { + let server = + only_server(r#"{"mcpServers": {"devtools": {"command": "npx", "args": ["-y", "chrome-devtools-mcp"]}}}"#); + + let McpServerConfig::Stdio(StdioServerConfig { command, args, defer_tools, .. }) = server else { + panic!("expected a stdio server, got {server:?}"); + }; + assert_eq!(command, "npx"); + assert_eq!(args, ["-y", "chrome-devtools-mcp"]); + assert!(!defer_tools.defers_tools()); +} + +#[test] +fn from_json_rejects_unknown_server_types() { + for server in [ + r#"{"type":"deferTools","servers":{}}"#, + r#"{"type":"htp","url":"https://example.com"}"#, + r#"{"type":"sse","url":"https://example.com"}"#, + r#"{"type":"unknown","servers":{"bad":{"type":"in-memory"}}}"#, + ] { + let json = format!(r#"{{"servers":{{"bad":{server}}}}}"#); + assert!(McpConfig::from_json(&json).is_err(), "unknown server type was accepted: {server}"); + } } -fn server_json(name: &str, body: &str) -> String { - format!(r#"{{ "servers": {{ "{name}": {body} }} }}"#) +#[test] +fn from_json_rejects_unknown_fields_for_all_transports() { + for server in [ + r#"{"command":"x","direct_tool":["bash"]}"#, + r#"{"type":"http","url":"https://example.com","direct_tool":["bash"]}"#, + r#"{"type":"in-memory","direct_tool":["bash"]}"#, + r#"{"command":"x","deferTools":true,"direct_tools":["bash"]}"#, + ] { + let json = format!(r#"{{"servers":{{"bad":{server}}}}}"#); + assert!(McpConfig::from_json(&json).is_err(), "unknown field was accepted: {server}"); + } } -macro_rules! with_env { - ([$( ($k:expr, $v:expr) ),+ $(,)?], $body:expr) => {{ - unsafe { $( env::set_var($k, $v); )+ } - let _result = $body; - unsafe { $( env::remove_var($k); )+ } - _result - }}; +#[test] +fn from_json_rejects_a_zero_oauth_callback_port() { + let json = r#"{"servers":{"bad":{"type":"http","url":"https://example.com/mcp","oauth":{"clientId":"client","callbackPort":0}}}}"#; + + assert!(McpConfig::from_json(json).is_err()); } -fn assert_http(server: McpServer, expected_name: &str, expected_url: &str) -> McpServer { - match &server.transport { - McpTransport::Http(c) => { - assert_eq!(server.name, expected_name); - assert_eq!(c.transport.uri.to_string(), expected_url); - } - other => panic!("Expected Http config, got {other:?}"), +#[test] +fn defer_tools_accepts_booleans_and_rules_for_all_transports() { + let config = McpConfig::from_json( + r#"{"servers":{"all":{"command":"a","deferTools":true},"stdio":{"command":"x","deferTools":{"include":["lsp_*"],"exclude":["lsp_rename"]}},"http":{"type":"http","url":"https://example.com","deferTools":{"exclude":["bash"]}},"memory":{"type":"in-memory","deferTools":{"include":["read"]}},"direct":{"type":"http","url":"https://example.com"}}}"#, + ) + .unwrap(); + + assert_eq!(config.servers["all"].defer_tools(), &ToolExposure::deferred_all()); + assert_eq!(config.servers["stdio"].defer_tools(), &deferred(&["lsp_*"], &["lsp_rename"])); + assert_eq!(config.servers["http"].defer_tools(), &deferred(&[], &["bash"])); + assert_eq!(config.servers["memory"].defer_tools(), &deferred(&["read"], &[])); + assert_eq!(config.servers["direct"].defer_tools(), &ToolExposure::ModelVisible); +} + +#[test] +fn defer_tools_rejects_unknown_rule_fields() { + for rules in [r#"{"allow":["lsp_*"]}"#, r#"{"deny":["bash"]}"#, r#"{"include":["lsp_*"],"direct":["bash"]}"#] { + let json = format!(r#"{{"servers":{{"coding":{{"command":"x","deferTools":{rules}}}}}}}"#); + assert!(McpConfig::from_json(&json).is_err(), "unknown deferTools field was accepted: {rules}"); } - server -} - -#[tokio::test] -async fn test_parse_stdio_config() { - let json = server_json( - "githubMcp", - r#"{ - "type": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-github"], - "env": { "GITHUB_TOKEN": "$GITHUB_TOKEN" } - }"#, - ); - with_env!([("GITHUB_TOKEN", "test_token")], { - let server = parse_one(&json); - assert!(!server.has_deferred_tools()); - match server.transport { - McpTransport::Stdio { command, args, env } => { - assert_eq!(server.name, "githubMcp"); - assert_eq!(command, "npx"); - assert_eq!(args, vec!["-y", "@modelcontextprotocol/server-github"]); - assert_eq!(env.get("GITHUB_TOKEN").unwrap(), "test_token"); - } - other => panic!("Expected Stdio config, got {other:?}"), - } - }); -} - -#[tokio::test] -async fn test_parse_http_oauth_config() { - let json = server_json( - "slack", - r#"{ - "type": "http", - "url": "https://mcp.slack.com/mcp", - "oauth": { - "clientId": "1601185624273.8899143856786", - "callbackPort": 3118 - } - }"#, +} + +#[test] +fn defer_tools_serializes_only_when_tools_are_deferred() { + let config = McpConfig::from_json( + r#"{"servers":{"coding":{"command":"x","deferTools":{"exclude":["bash","lsp_*"]}},"direct":{"command":"y","deferTools":false},"full":{"command":"z","deferTools":true}}}"#, + ) + .unwrap(); + + let value = serde_json::to_value(config).unwrap(); + + assert_eq!(value["servers"]["coding"]["deferTools"], serde_json::json!({"exclude":["bash", "lsp_*"]})); + assert!(value["servers"]["direct"].get("deferTools").is_none()); + assert_eq!(value["servers"]["full"]["deferTools"], serde_json::json!(true)); +} + +#[test] +fn deferred_tool_rules_partition_tools_with_exclude_winning() { + let server = only_server( + r#"{"servers":{"coding":{"command":"server","deferTools":{"include":["lsp_*","bash"],"exclude":["lsp_rename"]}}}}"#, ); + let exposure = server.defer_tools(); + + assert!(exposure.is_deferred(&tool("lsp_hover"))); + assert!(!exposure.is_deferred(&tool("lsp_rename"))); + assert!(exposure.is_deferred(&tool("bash"))); + assert!(!exposure.is_deferred(&tool("read_file"))); +} + +#[test] +fn deferred_tool_rules_match_tool_annotations() { + let server = only_server(r#"{"servers":{"coding":{"command":"x","deferTools":{"exclude":[{"readOnly":true}]}}}}"#); + let mut read_only = tool("read_file"); + read_only.annotations = Some(ToolAnnotations::new().read_only(true)); + + assert!(!server.defer_tools().is_deferred(&read_only)); + assert!(server.defer_tools().is_deferred(&tool("write_file"))); +} + +#[test] +fn defer_all_tools_defers_every_server_and_keeps_existing_rules() { + let mut config = McpConfig::from_json( + r#"{"servers":{"a":{"type":"stdio","command":"a"},"b":{"type":"http","url":"https://example.com"},"coding":{"command":"server","deferTools":{"exclude":["bash"]}}}}"#, + ) + .unwrap(); + + config.defer_all_tools(); + + assert!(config.servers.values().all(|server| server.defer_tools().defers_tools())); + let coding = config.servers["coding"].defer_tools(); + assert!(!coding.is_deferred(&tool("bash"))); + assert!(coding.is_deferred(&tool("read_file"))); +} + +#[test] +fn merging_no_configs_is_empty() { + assert!(std::iter::empty::().collect::().is_empty()); +} + +#[test] +fn merging_keeps_disjoint_servers() { + let merged: McpConfig = [ + config(r#"{"servers": {"alpha": {"type": "stdio", "command": "a"}}}"#), + config(r#"{"servers": {"beta": {"type": "stdio", "command": "b"}}}"#), + ] + .into_iter() + .collect(); + + assert_eq!(merged.servers.keys().collect::>(), ["alpha", "beta"]); +} + +#[test] +fn merging_lets_the_last_config_win_including_its_exposure() { + let first = + config(r#"{"servers":{"coding":{"type":"stdio","command":"first","deferTools":{"exclude":["bash"]}}}}"#); + let second = config(r#"{"servers":{"coding":{"type":"stdio","command":"second"}}}"#); + + assert_eq!(coding([first.clone(), second.clone()]), command("second", ToolExposure::ModelVisible)); + assert_eq!(coding([second, first]), command("first", deferred(&[], &["bash"]))); +} + +#[test] +fn a_later_config_can_make_a_deferred_server_model_visible() { + let mut lazy = config(r#"{"servers":{"coding":{"type":"stdio","command":"lazy"}}}"#); + lazy.defer_all_tools(); + let visible = config(r#"{"servers":{"coding":{"type":"stdio","command":"visible","deferTools":false}}}"#); + + assert_eq!(coding([lazy, visible]), command("visible", ToolExposure::ModelVisible)); +} + +#[test] +fn reading_a_config_file_reports_read_and_parse_errors() { + let dir = tempdir().unwrap(); + let invalid = write_config(dir.path(), "bad.json", "not valid json"); + + assert!(matches!(McpConfig::from_json_file(dir.path().join("missing.json")), Err(ParseError::IoError(_)))); + assert!(matches!(McpConfig::from_json_file(invalid), Err(ParseError::JsonError(_)))); + assert!(matches!(McpConfig::from_json("{"), Err(ParseError::JsonError(_)))); +} + +#[test] +fn stdio_servers_expand_variables_and_run_in_the_given_directory() { + let json = r#"{"servers":{"github":{"type":"stdio","command":"npx","args":["--root","${WORKSPACE}/src"],"env":{"GITHUB_TOKEN":"$GITHUB_TOKEN"}}}}"#; + let vars = vars(&[("WORKSPACE", "/workspace"), ("GITHUB_TOKEN", "test_token")]); + + let Transport::Stdio { command, args, env, cwd } = transport(json, &vars) else { panic!("expected stdio") }; - let server = parse_one(&json); - match server.transport { - McpTransport::Http(McpHttpConfig { oauth: Some(oauth), .. }) => { - assert_eq!(oauth.client_id.as_deref(), Some("1601185624273.8899143856786")); - assert_eq!(oauth.callback_port.map(NonZeroU16::get), Some(3118)); + assert_eq!(command, "npx"); + assert_eq!(args, ["--root", "/workspace/src"]); + assert_eq!(env["GITHUB_TOKEN"], "test_token"); + assert_eq!(cwd.as_deref(), Some(Path::new("/cwd"))); +} + +#[test] +fn missing_variables_fail_resolution() { + let json = r#"{"servers":{"test":{"type":"stdio","command":"$MISSING_VAR"}}}"#; + + let result = try_transport(json, &vars(&[])); + + assert!(matches!(result, Err(ParseError::VarError(VarError::NotFound(name))) if name == "MISSING_VAR")); +} + +#[test] +fn http_servers_connect_to_their_expanded_url() { + let vars = vars(&[("HOST", "localhost"), ("PORT", "8080")]); + let json = r#"{"servers":{"remote":{"type":"http","url":"http://${HOST}:${PORT}/mcp"}}}"#; + + let Transport::Http { url, .. } = transport(json, &vars) else { panic!("expected http") }; + + assert_eq!(url, "http://localhost:8080/mcp"); +} + +#[test] +fn http_headers_expand_variables_and_are_marked_sensitive() { + let json = r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"Authorization":"Bearer ${TOKEN}"}}}}"#; + + let Transport::Http { headers, .. } = transport(json, &vars(&[("TOKEN", "expanded-token")])) else { + panic!("expected http") + }; + + assert_eq!(headers[AUTHORIZATION], "Bearer expanded-token"); + assert!(headers[AUTHORIZATION].is_sensitive()); +} + +#[test] +fn http_headers_keep_authorization_schemes_and_accept_any_header_casing() { + let vars = vars(&[("TOKEN", "secret"), ("API_KEY", "custom-secret")]); + for scheme in ["Sentry-Bearer", "Bearer", "bearer", "Basic", "Token"] { + for header_name in ["Authorization", "authorization", "AUTHORIZATION"] { + let json = format!( + r#"{{"servers":{{"remote":{{"type":"http","url":"http://localhost/mcp","headers":{{"{header_name}":"{scheme} $TOKEN","X-API-Key":"$API_KEY"}}}}}}}}"# + ); + + let Transport::Http { headers, .. } = transport(&json, &vars) else { panic!("expected http") }; + + assert_eq!(headers.get_all(AUTHORIZATION).iter().count(), 1); + assert_eq!(headers[AUTHORIZATION], format!("{scheme} secret")); + assert_eq!(headers["x-api-key"], "custom-secret"); } - other => panic!("Expected HTTP OAuth config, got {other:?}"), } } #[test] -fn test_rejects_zero_oauth_callback_port() { - let json = server_json( - "bad", - r#"{ - "type": "http", - "url": "https://example.com/mcp", - "oauth": { "clientId": "client", "callbackPort": 0 } - }"#, - ); +fn invalid_http_headers_fail_resolution() { + let invalid_value = + r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"X-Key":"bad\nvalue"}}}}"#; + let invalid_name = + r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"Bad Key":"value"}}}}"#; - assert!(McpConfig::from_json(&json).is_err()); + assert!(matches!( + try_transport(invalid_value, &vars(&[])), + Err(ParseError::InvalidHeaderValue { name, .. }) if name == "X-Key" + )); + assert!(matches!( + try_transport(invalid_name, &vars(&[])), + Err(ParseError::InvalidHeaderName { name, .. }) if name == "Bad Key" + )); } -#[tokio::test] -async fn test_parse_http_and_sse_configs() { - let json = server_json( - "mcpMesh", - r#"{ - "type": "http", - "url": "http://localhost:3000/mcp", - "headers": { "Authorization": "Bearer $API_TOKEN" } - }"#, - ); - let cfg = with_env!( - [("API_TOKEN", "secret_token")], - assert_http(parse_one(&json), "mcpMesh", "http://localhost:3000/mcp") - ); - if let McpTransport::Http(c) = cfg.transport { - assert!(c.transport.auth_header.is_none()); - assert_eq!(c.transport.custom_headers[&AUTHORIZATION], "Bearer secret_token"); +#[test] +fn http_servers_sign_in_with_aether_client_metadata_by_default() { + let oauth = http_oauth(None, HeaderMap::new()).expect("OAuth is offered"); + + assert!(matches!( + &oauth.registration, + OAuthClientRegistration::ClientMetadata(url) if url == "https://aether-agent.io/oauth/client-metadata.json" + )); + assert_eq!(oauth.callback_port, 3118); +} + +#[test] +fn http_servers_prefer_a_preregistered_client_id_over_client_metadata() { + let config = McpOAuthConfig { + client_id: Some("registered".to_string()), + client_metadata_url: Some("https://client.example/oauth/client.json".to_string()), + callback_port: NonZeroU16::new(4000), + }; + + let oauth = http_oauth(Some(config), HeaderMap::new()).expect("OAuth is offered"); + + assert!(matches!(&oauth.registration, OAuthClientRegistration::PreRegistered(id) if id == "registered")); + assert_eq!(oauth.callback_port, 4000); +} + +#[test] +fn http_servers_with_an_authorization_header_never_sign_in() { + let headers = HeaderMap::from_iter([(AUTHORIZATION, HeaderValue::from_static("Bearer static"))]); + + assert!(http_oauth(None, headers).is_none()); +} + +#[test] +fn oauth_settings_expand_variables() { + let vars = vars(&[("CLIENT_ID", "registered"), ("CIMD_URL", "https://client.example/oauth/client.json")]); + let oauth = |settings: &str| { + let json = format!( + r#"{{"servers":{{"linear":{{"type":"http","url":"https://example.com/mcp","oauth":{settings}}}}}}}"# + ); + let Transport::Http { oauth: Some(oauth), .. } = transport(&json, &vars) else { + panic!("expected HTTP transport with OAuth settings") + }; + oauth + }; + + let registered = oauth(r#"{"clientId":"$CLIENT_ID","callbackPort":4000}"#); + let metadata = oauth(r#"{"clientMetadataUrl":"$CIMD_URL"}"#); + + assert!(matches!(®istered.registration, OAuthClientRegistration::PreRegistered(id) if id == "registered")); + assert_eq!(registered.callback_port, 4000); + assert!(matches!( + &metadata.registration, + OAuthClientRegistration::ClientMetadata(url) if url == "https://client.example/oauth/client.json" + )); +} + +fn vars(pairs: &[(&str, &str)]) -> Vars { + pairs.iter().fold(Vars::new().with_env_lookup(|_| None), |vars, (name, value)| vars.with(*name, *value)) +} + +fn only_server(json: &str) -> McpServerConfig { + let mut servers = McpConfig::from_json(json).unwrap().servers; + assert_eq!(servers.len(), 1); + servers.pop_first().expect("one server").1 +} + +fn try_transport(json: &str, vars: &Vars) -> Result { + match only_server(json) { + McpServerConfig::Stdio(config) => config.into_transport(vars, Path::new("/cwd")), + McpServerConfig::Remote(config) => config.into_transport(vars), + McpServerConfig::InMemory(_) => panic!("in-memory servers have no transport"), } +} + +fn transport(json: &str, vars: &Vars) -> Transport { + try_transport(json, vars).unwrap() +} - let json = server_json("sseServer", r#"{ "type": "sse", "url": "http://localhost:4000/sse", "headers": {} }"#); - assert_http(parse_one(&json), "sseServer", "http://localhost:4000/sse"); +fn config(json: &str) -> McpConfig { + McpConfig::from_json(json).unwrap() } -#[tokio::test] -async fn test_missing_env_var_error() { - let json = server_json("test", r#"{ "type": "stdio", "command": "$MISSING_VAR", "args": [] }"#); - match parse_servers(&json).unwrap_err() { - ParseError::VarError(_) => (), - other => panic!("Expected VarError, got {other:?}"), +fn coding(configs: impl IntoIterator) -> (String, ToolExposure) { + match configs.into_iter().collect::().servers.remove("coding") { + Some(McpServerConfig::Stdio(StdioServerConfig { command, defer_tools, .. })) => (command, defer_tools), + other => panic!("expected a stdio server, got {other:?}"), } } -#[tokio::test] -async fn test_in_memory_config_is_declarative() { - let json = server_json( - "test", - r#"{ "type": "in-memory", "args": ["--root", "${WORKSPACE}"], "input": {"enabled": true} }"#, - ); - let vars = Vars::new().with("WORKSPACE", "/workspace"); - let mut servers = McpConfig::from_json(&json).unwrap().into_servers(&vars).unwrap(); - let McpTransport::InMemory { spec } = servers.remove(0).transport else { - panic!("expected in-memory transport"); +fn command(command: &str, exposure: ToolExposure) -> (String, ToolExposure) { + (command.to_string(), exposure) +} + +fn http_oauth(oauth: Option, headers: HeaderMap) -> Option { + let Transport::Http { oauth, .. } = http_transport("https://example.com/mcp".to_string(), headers, oauth) else { + panic!("expected http") }; - assert_eq!(spec.factory, "test"); - assert_eq!(spec.args, ["--root", "/workspace"]); - assert_eq!(spec.input, Some(serde_json::json!({"enabled": true}))); -} - -#[tokio::test] -async fn test_multiple_servers() { - let json = r#"{ - "servers": { - "server1": { "type": "stdio", "command": "node", "args": ["server.js"] }, - "server2": { - "type": "http", - "url": "http://localhost:3000/mcp", - "headers": { "Authorization": "$TOKEN" } - } - } - }"#; - with_env!([("TOKEN", "test")], { - assert_eq!(parse_servers(json).unwrap().len(), 2); - }); -} - -#[tokio::test] -async fn test_env_var_in_url() { - let json = server_json("test", r#"{ "type": "http", "url": "http://${HOST}:${PORT}/mcp" }"#); - with_env!([("HOST", "localhost"), ("PORT", "8080")], { - assert_http(parse_one(&json), "test", "http://localhost:8080/mcp"); - }); -} - -#[tokio::test] -async fn test_parse_per_server_defer_tools_config() { - let json = r#"{ - "servers": { - "github": { - "type": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-github"], - "deferTools": true - }, - "sentry": { "type": "http", "url": "https://sentry.example.com/mcp" } - } - }"#; - let servers = parse_servers(json).unwrap(); - assert_eq!(servers.len(), 2); - assert!(servers.iter().find(|s| s.name == "github").unwrap().has_deferred_tools()); - assert!(!servers.iter().find(|s| s.name == "sentry").unwrap().has_deferred_tools()); + oauth } -#[test] -fn test_rejects_unknown_server_type() { - let json = server_json("outer", r#"{ "type": "unknown", "servers": { "bad": { "type": "in-memory" } } }"#); - assert!(McpConfig::from_json(&json).is_err()); +fn deferred(allow: &[&str], deny: &[&str]) -> ToolExposure { + let matchers = |patterns: &[&str]| patterns.iter().copied().map(ToolMatcher::name).collect(); + ToolExposure::Deferred(ToolFilter { allow: matchers(allow), deny: matchers(deny) }) +} + +fn tool(name: &str) -> Tool { + Tool::new(name.to_string(), "", Arc::new(serde_json::Map::new())) +} + +fn write_config(dir: &Path, name: &str, json: &str) -> PathBuf { + let path = dir.join(name); + std::fs::write(&path, json).unwrap(); + path } diff --git a/crates/aether-core/tests/mcp/instructions_tests.rs b/crates/aether-core/tests/mcp/instructions_tests.rs index 041b03ce9..9452332ed 100644 --- a/crates/aether-core/tests/mcp/instructions_tests.rs +++ b/crates/aether-core/tests/mcp/instructions_tests.rs @@ -1,76 +1,52 @@ -use aether_core::core::Prompt; -use aether_core::testing::{FakeMcpServer, McpTestBuilder, mcp_instructions as instructions, test_agent}; +use aether_core::core::{McpInstructions, Prompt}; +use aether_core::mcp::mcp; +use aether_core::testing::{mcp_instructions, ready_catalog, test_agent}; use llm::ChatMessage; use llm::testing::llm_response; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::testing::FakeMcpServer; use std::error::Error; #[tokio::test] -async fn test_fake_mcp_server_has_instructions() { - let mcp_test = McpTestBuilder::new().server("test", FakeMcpServer::new()).build().await; - let instructions = mcp_test.snapshot().model_instructions(); +async fn each_server_contributes_its_own_instructions() { + let servers = vec![ + ServerSpec::in_process("server1", FakeMcpServer::math().with_instructions("First server")), + ServerSpec::in_process("server2", FakeMcpServer::math().with_instructions("Second server")), + ServerSpec::in_process("silent", FakeMcpServer::math()), + ]; + let instructions = ready_catalog(mcp("/workspace").with_servers(servers)).await.instructions(); - // FakeMcpServer does provide instructions, so we should get them - assert_eq!(instructions.len(), 1); - assert!(instructions.get("test").unwrap().contains("A fake MCP server for testing")); + assert_eq!(instructions, mcp_instructions(&[("server1", "First server"), ("server2", "Second server")])); } #[tokio::test] -async fn test_multiple_servers_with_instructions() { - let mcp_test = McpTestBuilder::new() - .server("server1", FakeMcpServer::new()) - .server("server2", FakeMcpServer::new()) - .build() - .await; - let instructions = mcp_test.snapshot().model_instructions(); +async fn progressive_discovery_renders_apart_from_the_server_instructions() { + let only_deferred = McpInstructions { progressive_discovery: true, ..McpInstructions::default() }; - // Both servers should have instructions - assert_eq!(instructions.len(), 2); - assert!(instructions.get("server1").unwrap().contains("A fake MCP server for testing")); - assert!(instructions.get("server2").unwrap().contains("A fake MCP server for testing")); -} - -#[tokio::test] -async fn test_format_mcp_instructions_xml_structure() { - let formatted = Prompt::McpInstructions(instructions(&[("coding", "Use absolute paths.")])).build().await.unwrap(); + let formatted = Prompt::McpInstructions(only_deferred).build().await.unwrap(); - // Check for XML tags with server names - assert!(formatted.contains("")); - assert!(formatted.contains("\n")); - assert!(formatted.contains("Use absolute paths.")); - assert!(formatted.contains("# MCP Server Instructions")); + assert!(formatted.contains("aether mcp")); + assert!(!formatted.contains("")); - assert!(formatted.contains("")); - assert!(formatted.contains("Use absolute paths.")); - assert!(formatted.contains("Always confirm.")); -} + let formatted = Prompt::McpInstructions(instructions).build().await.unwrap(); -#[tokio::test] -async fn test_format_mcp_instructions_is_deterministically_ordered() { - let a = Prompt::McpInstructions(instructions(&[("zebra", "Z"), ("alpha", "A")])).build().await.unwrap(); - let b = Prompt::McpInstructions(instructions(&[("alpha", "A"), ("zebra", "Z")])).build().await.unwrap(); - assert_eq!(a, b); - let alpha_pos = a.find("name=\"alpha\"").unwrap(); - let zebra_pos = a.find("name=\"zebra\"").unwrap(); - assert!(alpha_pos < zebra_pos, "alphabetical order: alpha before zebra"); + assert!(formatted.starts_with("# MCP Server Instructions"), "{formatted}"); + let alpha = formatted.find("\nUse absolute paths.\n").expect("alpha"); + let zebra = formatted.find("\nAlways confirm.\n").expect("zebra"); + assert!(alpha < zebra); } #[tokio::test] async fn test_agent_builder_includes_mcp_instructions_in_system_prompt() -> Result<(), Box> { let result = test_agent() + .mcp_server("test-server", FakeMcpServer::math().with_instructions("Use the math tools")) .llm_responses(&[llm_response().text(&["done"]).build()]) .system_prompt(Prompt::text("You are a test agent")) - .system_prompt(Prompt::McpInstructions(instructions(&[("test-server", "Test instructions")]))) .user_text("test") .run_with_context() .await?; @@ -80,14 +56,15 @@ async fn test_agent_builder_includes_mcp_instructions_in_system_prompt() -> Resu let ChatMessage::System { content, .. } = first_message else { panic!("Expected system message, got: {first_message:?}"); }; - assert!(content.contains("")); - assert!(content.contains("Test instructions")); + assert!(content.contains("You are a test agent")); + assert!(content.contains("\nUse the math tools\n")); Ok(()) } #[tokio::test] async fn test_agent_builder_works_without_mcp_instructions() -> Result<(), Box> { let result = test_agent() + .without_mcp() .llm_responses(&[llm_response().text(&["done"]).build()]) .system_prompt(Prompt::text("You are a test agent")) .user_text("test") diff --git a/crates/aether-core/tests/mcp/ipc_tests.rs b/crates/aether-core/tests/mcp/ipc_tests.rs index f84758d4e..a5504df70 100644 --- a/crates/aether-core/tests/mcp/ipc_tests.rs +++ b/crates/aether-core/tests/mcp/ipc_tests.rs @@ -1,116 +1,177 @@ -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, McpTestBuilder}; -use mcp_utils::ServiceExt; -use mcp_utils::client::{DeferredToolRules, ToolExposure, ToolFilter, ToolMatcher}; -use mcp_utils::tool_gateway::{LIST_SERVERS_TOOL, connect}; -use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, CreateTaskResult, DetailedTask, Task, TaskPayload, - TaskStatus, +use aether_core::mcp::{ + DeferredServer, DeferredTool, DeferredToolsClient, McpBuilder, McpCatalog, ToolExposure, ToolFilter, ToolMatcher, + mcp, }; -use serde_json::json; +use aether_core::testing::ready_catalog; +use futures::StreamExt; +use mcp_utils::client::{ToolCallError, ToolCallEvent}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::testing::{FakeMcpServer, FakeTask, FakeTool, FakeToolResponse, args, completed_task_payload, task}; +use rmcp::model::{CallToolResult, DetailedTask, ErrorCode, TaskStatus}; +use rmcp::service::ServiceError; +use serde_json::{Map, json}; +use std::time::Duration; #[tokio::test] async fn no_deferred_servers_do_not_create_a_gateway() { - let mcp_test = McpTestBuilder::new().server("math", FakeMcpServer::new()).build().await; + let mcp = + ready_catalog(mcp("/workspace").with_servers(vec![ServerSpec::in_process("math", FakeMcpServer::math())])) + .await; - assert!(mcp_test.gateway_endpoint().is_none()); + assert!(mcp.deferred_tools_socket().is_none()); } #[tokio::test] async fn gateway_discovers_and_calls_only_deferred_tools() { - let exposure = ToolExposure::Deferred(DeferredToolRules::new(&["divide_numbers"], &[])); - let mcp_test = McpTestBuilder::new().server_with_exposure("math", FakeMcpServer::new(), exposure).build().await; - let client = - ().serve(connect(mcp_test.gateway_endpoint().expect("deferred gateway exists")).await.unwrap()).await.unwrap(); + let exposure = + ToolExposure::Deferred(ToolFilter { allow: vec![ToolMatcher::name("divide_numbers")], deny: Vec::new() }); + let mcp = + ready_catalog(mcp("/workspace").with_server(ServerSpec::in_process("math", FakeMcpServer::math()), exposure)) + .await; + let client = connect(&mcp).await; - let servers = client.call_tool_once(CallToolRequestParams::new(LIST_SERVERS_TOOL)).await.unwrap(); - let CallToolResponse::Complete(servers) = servers else { panic!("server discovery is complete") }; assert_eq!( - servers.structured_content, - Some(json!([{"name": "math", "description": "A fake MCP server for testing"}])) + client.servers().await.unwrap(), + [DeferredServer { + name: "math".to_string(), + description: "A fake MCP server for testing".to_string(), + tools: vec![DeferredTool { + name: "divide_numbers".to_string(), + description: Some("Divides two numbers".to_string()) + }], + }] ); - - let tools = client.list_tools(None).await.unwrap().tools; - let names = tools.iter().map(|tool| tool.name.as_ref()).collect::>(); - assert!(names.contains(&"math__divide_numbers")); - assert!(!names.contains(&"math__add_numbers")); - let divide = tools.iter().find(|tool| tool.name == "math__divide_numbers").unwrap(); + let divide = client.tool("math", "divide_numbers").await.unwrap().expect("divide_numbers is deferred"); assert_eq!(divide.description.as_deref(), Some("Divides two numbers")); + assert!(client.tool("math", "add_numbers").await.unwrap().is_none()); - let result = client - .call_tool_once( - CallToolRequestParams::new("math__divide_numbers") - .with_arguments(json!({"a": 8, "b": 2}).as_object().unwrap().clone()), - ) - .await - .unwrap(); - let CallToolResponse::Complete(result) = result else { - panic!("gateway returns a complete result"); - }; + let result = call(&client, "math", "divide_numbers", json!({"a": 8, "b": 2})).await.unwrap(); assert_eq!(result.structured_content, Some(json!({"quotient": 4}))); - let error = client - .call_tool_once( - CallToolRequestParams::new("math__add_numbers") - .with_arguments(json!({"a": 1, "b": 2}).as_object().unwrap().clone()), - ) + let error = call(&client, "math", "add_numbers", json!({"a": 1, "b": 2})) .await .expect_err("model-visible tools are rejected by the deferred route"); - assert!(error.to_string().contains("exposed directly") || error.to_string().contains("Tool not found")); + assert!(error.to_string().contains("exposed directly; call 'math__add_numbers'")); } #[tokio::test] async fn gateway_discovery_and_execution_share_the_tool_filter() { let server = FakeMcpServer::new().with_tool(FakeTool::new("secret").responds(FakeToolResponse::text("hidden"))); let filter = ToolFilter { allow: Vec::new(), deny: vec![ToolMatcher::name("math__secret")] }; - let mcp_test = McpTestBuilder::new().deferred_server("math", server).tool_filter(filter).build().await; - let client = - ().serve(connect(mcp_test.gateway_endpoint().expect("deferred gateway exists")).await.unwrap()).await.unwrap(); - - let tools = client.list_tools(None).await.unwrap().tools; - assert!(!tools.iter().any(|tool| tool.name == "math__secret")); - let error = client - .call_tool_once(CallToolRequestParams::new("math__secret")) - .await - .expect_err("filtered tools cannot be executed"); - assert!(error.to_string().contains("Tool not found")); + let mcp = ready_catalog(deferred("math", server).with_tool_filter(filter)).await; + let client = connect(&mcp).await; + + assert!(client.tool("math", "secret").await.unwrap().is_none()); + let error = call(&client, "math", "secret", json!({})).await.expect_err("filtered tools cannot be executed"); + assert!(matches!( + error, + ToolCallError::Call(ServiceError::McpError(error)) + if error.code == ErrorCode::INVALID_PARAMS && error.message.contains("Tool not found") + )); } #[tokio::test] async fn gateway_reduces_mcp_task_completion_to_the_final_result() { - let now = chrono::Utc::now().to_rfc3339(); - let working = Task::new("gateway-task", TaskStatus::Working, now.clone(), now.clone()).with_poll_interval_ms(1); - let completed = Task::new("gateway-task", TaskStatus::Completed, now.clone(), now); - let final_result = CallToolResult::structured(json!({"done": true})); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("background").responds(FakeToolResponse::task(CreateTaskResult::new(working)))) - .with_task( - "gateway-task", - [DetailedTask::new( - completed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(final_result).unwrap()).unwrap(), - }, - )], - ); - let mcp_test = McpTestBuilder::new().deferred_server("math", server).build().await; - let client = - ().serve(connect(mcp_test.gateway_endpoint().expect("deferred gateway exists")).await.unwrap()).await.unwrap(); - - let response = client.call_tool_once(CallToolRequestParams::new("math__background")).await.unwrap(); - - let CallToolResponse::Complete(result) = response else { panic!("gateway resolves MCP Tasks") }; - assert_eq!(result.structured_content, Some(json!({"done": true}))); + let payload = completed_task_payload(CallToolResult::structured(json!({"done": true}))); + let completed = DetailedTask::new(task("gateway-task", TaskStatus::Completed), payload); + let mcp = + ready_catalog(deferred("math", FakeTask::new("gateway-task").states([completed]).server("background"))).await; + let client = connect(&mcp).await; + + let events = client.call("math", "background", Map::new(), TIMEOUT).collect::>().await; + + assert!( + !events.iter().any(|event| matches!(event, ToolCallEvent::TaskCreated(_))), + "the gateway resolves MCP Tasks itself" + ); + assert!(matches!( + events.last(), + Some(ToolCallEvent::Done(Ok(result))) if result.structured_content == Some(json!({"done": true})) + )); } #[tokio::test] -async fn gateway_socket_is_removed_with_the_runtime() { - let mcp_test = McpTestBuilder::new().deferred_server("math", FakeMcpServer::new()).build().await; - let endpoint = mcp_test.gateway_endpoint().expect("deferred gateway exists").to_path_buf(); +async fn the_socket_relays_upstream_progress() { + let ci = FakeMcpServer::new() + .with_tool(FakeTool::new("build").responds(FakeToolResponse::text("built").progress_message(1.0, "compiling"))); + let mcp = ready_catalog(deferred("ci", ci)).await; + let client = connect(&mcp).await; + + let events = client.call("ci", "build", Map::new(), TIMEOUT).collect::>().await; + + assert!(events.iter().any(|event| matches!( + event, + ToolCallEvent::Progress(progress) if progress.message.as_deref() == Some("compiling") + ))); + assert!(matches!(events.last(), Some(ToolCallEvent::Done(Ok(_))))); +} + +#[tokio::test] +async fn listing_servers_names_only_servers_with_deferred_tools_and_their_deferred_tools() { + let partial = ToolExposure::Deferred(ToolFilter { allow: vec![ToolMatcher::name("add_*")], deny: vec![] }); + let mcp = ready_catalog( + deferred("lazy", FakeMcpServer::math()) + .with_server(ServerSpec::in_process("partial", FakeMcpServer::math()), partial) + .with_servers(vec![ServerSpec::in_process("direct", FakeMcpServer::math())]), + ) + .await; + let client = connect(&mcp).await; + + let servers = client.servers().await.unwrap(); + + let tools = |server: &DeferredServer| server.tools.iter().map(|tool| tool.name.clone()).collect::>(); + assert_eq!(servers.iter().map(|server| server.name.as_str()).collect::>(), ["lazy", "partial"]); + assert_eq!(tools(&servers[0]), ["add_numbers", "divide_numbers"]); + assert_eq!(tools(&servers[1]), ["add_numbers"]); + assert_eq!(servers[1].tools[0].description.as_deref(), Some("Adds two numbers together")); +} + +#[tokio::test] +async fn disconnecting_from_the_socket_cancels_the_server_task() { + let server = FakeTask::new("task-1").server("deferred"); + let state = server.state(); + let mcp = ready_catalog(deferred("lazy", server)).await; + let client = connect(&mcp).await; + let in_flight = tokio::spawn(client.call("lazy", "deferred", Map::new(), TIMEOUT).result()); + state.wait_until(|state| !state.task_get_ids().is_empty()).await; + + client.close().await; + + assert!(in_flight.await.unwrap().is_err()); + state.wait_until(|state| state.task_cancel_ids() == ["task-1"]).await; +} + +#[tokio::test] +async fn the_socket_is_removed_once_the_session_is_dropped() { + let mcp = ready_catalog(deferred("math", FakeMcpServer::math())).await; + let endpoint = mcp.deferred_tools_socket().expect("deferred gateway exists").to_path_buf(); let directory = endpoint.parent().unwrap().to_path_buf(); + let mut changes = mcp.changes(); assert!(endpoint.exists()); - drop(mcp_test); + drop(mcp); + while changes.next().await.is_some() {} assert!(!endpoint.exists()); assert!(!directory.exists()); } + +const TIMEOUT: Duration = Duration::from_mins(1); + +fn deferred(name: &str, server: FakeMcpServer) -> McpBuilder { + mcp("/workspace").with_server(ServerSpec::in_process(name, server), ToolExposure::deferred_all()) +} + +async fn connect(mcp: &McpCatalog) -> DeferredToolsClient { + let socket = mcp.deferred_tools_socket().expect("deferred gateway exists").to_path_buf(); + DeferredToolsClient::connect(socket).await.unwrap() +} + +async fn call( + client: &DeferredToolsClient, + server: &str, + tool: &str, + arguments: serde_json::Value, +) -> Result { + client.call(server, tool, args(arguments), TIMEOUT).result().await +} diff --git a/crates/aether-core/tests/mcp/mrtr_tests.rs b/crates/aether-core/tests/mcp/mrtr_tests.rs index db5d4f73f..dfc193117 100644 --- a/crates/aether-core/tests/mcp/mrtr_tests.rs +++ b/crates/aether-core/tests/mcp/mrtr_tests.rs @@ -1,130 +1,21 @@ -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, McpTestBuilder}; -use rmcp::model::{ - CallToolResponse, ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction, InputRequest, InputRequests, - InputRequiredResult, -}; +use super::support::call_tool; +use aether_core::testing::test_agent; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse, accept_with_name, name_input}; use serde_json::json; -const TRICKY_STATE: &str = "st.ate/with+special=chars and spaces \"quotes\"\n\ttab"; - #[tokio::test] -async fn form_round_trip_forwards_response_and_echoes_state() { +async fn input_requests_are_elicited_from_the_host_and_answered() { let server = FakeMcpServer::new().with_tool( - FakeTool::new("form") - .responds(input_required("answer", form_request("Name?"), TRICKY_STATE)) - .when_state(TRICKY_STATE, FakeToolResponse::text("done")), + FakeTool::new("form").responds(name_input("asked")).when_state("asked", FakeToolResponse::text("done")), ); let state = server.state(); - let test = - McpTestBuilder::new().server("mrtr", server).elicitation_response(accept_with_name("Ferris")).build().await; - - let result = test.call("mrtr", "form", json!({})).await.result.expect("form round trip completes"); - - assert!(result.result.contains("done")); - let calls = state.calls_for("form"); - assert_eq!(calls.len(), 2); - assert_eq!(calls[1].request.request_state.as_deref(), Some(TRICKY_STATE)); - assert_eq!(calls[1].request.input_responses.as_ref().unwrap()["answer"]["action"], "accept"); - assert_eq!(calls[1].request.input_responses.as_ref().unwrap()["answer"]["content"]["name"], "Ferris"); - let elicitations = test.elicitations(); - assert_eq!(elicitations.len(), 1); - assert_eq!(elicitations[0].server_name, "mrtr"); - assert!(matches!( - &elicitations[0].request, - ElicitRequestParams::FormElicitationParams { message, .. } if message == "Name?" - )); -} + let agent = test_agent().mcp_server("mrtr", server).elicitation_response(accept_with_name("Ferris")); -#[tokio::test] -async fn url_elicitation_and_state_only_polling_complete_without_reprompting() { - let server = FakeMcpServer::new().with_tool( - FakeTool::new("url_then_poll") - .responds(input_required( - "consent", - url_request("Authorize to continue", "https://example.com/auth", "el-1"), - "url-opened", - )) - .when_state("url-opened", InputRequiredResult::from_request_state("poll-1")) - .when_state("poll-1", FakeToolResponse::text("url-done")), - ); - let state = server.state(); - let test = McpTestBuilder::new() - .server("mrtr", server) - .elicitation_response(ElicitResult::new(ElicitationAction::Accept)) - .build() - .await; - - let result = test.call("mrtr", "url_then_poll", json!({})).await.result.expect("URL flow completes"); - - assert!(result.result.contains("url-done")); - assert_eq!(state.calls_for("url_then_poll").len(), 3); - let elicitations = test.elicitations(); - assert_eq!(elicitations.len(), 1, "state-only polling must not prompt again"); - assert!(matches!( - &elicitations[0].request, - ElicitRequestParams::UrlElicitationParams { url, elicitation_id, .. } - if url == "https://example.com/auth" && elicitation_id == "el-1" - )); -} - -#[tokio::test] -async fn progress_notifications_are_forwarded_across_rounds() { - let server = FakeMcpServer::new().with_tool( - FakeTool::new("progress_rounds") - .responds( - FakeToolResponse::new(input_required("answer", form_request("Name?"), "progress")) - .progress(1.0, Some(2.0)), - ) - .when_state("progress", FakeToolResponse::text("progress-done").progress(2.0, Some(2.0))), - ); - let test = - McpTestBuilder::new().server("mrtr", server).elicitation_response(accept_with_name("Ferris")).build().await; - - let outcome = test.call("mrtr", "progress_rounds", json!({})).await; - - assert!(outcome.result.expect("progress tool completes").result.contains("progress-done")); - assert_eq!(outcome.progress.iter().map(|event| event.progress).collect::>(), vec![1.0, 2.0]); -} - -#[tokio::test] -async fn invalid_input_required_result_becomes_a_tool_error() { - let server = - FakeMcpServer::new().with_tool(FakeTool::new("invalid").responds(InputRequiredResult::new(None, None))); - let test = McpTestBuilder::new().server("mrtr", server).build().await; - - let error = test.call("mrtr", "invalid", json!({})).await.result.expect_err("invalid result fails"); - - assert!(error.error.contains("without any input requests or request state"), "{}", error.error); -} - -fn form_request(message: &str) -> InputRequest { - InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::FormElicitationParams { - meta: None, - message: message.to_string(), - requested_schema: serde_json::from_value(json!({ - "type": "object", - "properties": { "name": { "type": "string" } }, - "required": ["name"] - })) - .unwrap(), - })) -} - -fn url_request(message: &str, url: &str, elicitation_id: &str) -> InputRequest { - InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::UrlElicitationParams { - meta: None, - message: message.to_string(), - url: url.to_string(), - elicitation_id: elicitation_id.to_string(), - })) -} - -fn input_required(key: &str, request: InputRequest, state: &str) -> CallToolResponse { - let mut requests = InputRequests::new(); - requests.insert(key.to_string(), request); - InputRequiredResult::new(Some(requests), Some(state.to_string())).into() -} + let run = call_tool(agent, "mrtr__form", json!({})).await; -fn accept_with_name(name: &str) -> ElicitResult { - ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": name })) + assert!(run.result().expect("the input round trip completes").result.contains("done")); + let answers = state.calls_for("form")[1].request.input_responses.clone().expect("the tool was answered"); + assert_eq!(answers["name"]["content"]["name"], "Ferris"); + assert_eq!(run.elicitations.len(), 1); + assert_eq!(run.elicitations[0].server_name, "mrtr"); } diff --git a/crates/aether-core/tests/mcp/oauth_tests.rs b/crates/aether-core/tests/mcp/oauth_tests.rs deleted file mode 100644 index 154fb09de..000000000 --- a/crates/aether-core/tests/mcp/oauth_tests.rs +++ /dev/null @@ -1,317 +0,0 @@ -use aether_auth::{OAuthError, OAuthHandler, accept_oauth_callback}; -use aether_core::mcp::mcp; -use aether_core::testing::{FakeMcpServer, fake_mcp}; -use futures::future::BoxFuture; -use mcp_utils::client::{ - DeferredToolRules, ElicitingOAuthHandler, McpClientEvent, McpManager, OAuthHandlerContext, OAuthHandlerFactory, - RuntimeMcpServer, RuntimeMcpTransport, ToolExposure, -}; -use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; -use std::sync::Arc; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::sync::mpsc; -use tokio::task::JoinHandle; -use utils::mcp_status::{McpServerAuthCapability, McpServerStatus}; - -struct FailingHttpEndpoint { - uri: String, - task: JoinHandle<()>, -} - -impl FailingHttpEndpoint { - async fn bind() -> Self { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let uri = format!("http://{}/mcp", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - while let Ok((stream, _)) = listener.accept().await { - drop(stream); - } - }); - Self { uri, task } - } - - fn server(&self, name: &str, exposure: ToolExposure) -> RuntimeMcpServer { - RuntimeMcpServer::new( - name, - RuntimeMcpTransport::Http(StreamableHttpClientTransportConfig::with_uri(self.uri.as_str()).into()), - exposure, - ) - } -} - -impl Drop for FailingHttpEndpoint { - fn drop(&mut self) { - self.task.abort(); - } -} - -struct UnauthorizedHttpEndpoint { - uri: String, - task: JoinHandle<()>, -} - -impl UnauthorizedHttpEndpoint { - async fn bind(challenge: &'static str) -> Self { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let uri = format!("http://{}/mcp", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - while let Ok((mut stream, _)) = listener.accept().await { - let mut request = [0; 4096]; - let _ = stream.read(&mut request).await; - let response = format!( - "HTTP/1.1 401 Unauthorized\r\nWWW-Authenticate: {challenge}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ); - let _ = stream.write_all(response.as_bytes()).await; - } - }); - Self { uri, task } - } - - fn server(&self, name: &str) -> RuntimeMcpServer { - RuntimeMcpServer::new( - name, - RuntimeMcpTransport::Http(StreamableHttpClientTransportConfig::with_uri(self.uri.as_str()).into()), - ToolExposure::ModelVisible, - ) - } -} - -impl Drop for UnauthorizedHttpEndpoint { - fn drop(&mut self) { - self.task.abort(); - } -} - -fn test_manager(with_oauth: bool) -> McpManager { - let (event_tx, _) = mpsc::channel::(50); - let factory = if with_oauth { Some(fake_oauth_handler_factory()) } else { None }; - McpManager::new(event_tx, factory) -} - -struct CancellingOAuthHandler; - -impl OAuthHandler for CancellingOAuthHandler { - fn redirect_uri(&self) -> &'static str { - "http://127.0.0.1:0/oauth2callback" - } - - fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { - Box::pin(async { Err(OAuthError::UserCancelled) }) - } -} - -fn fake_oauth_handler_factory() -> OAuthHandlerFactory { - Arc::new(|_ctx| Ok(Arc::new(CancellingOAuthHandler))) -} - -#[tokio::test] -async fn cancelling_handler_returns_user_cancelled() { - let handler = CancellingOAuthHandler; - let result = handler.authorize("https://example.com/auth").await; - assert!(matches!(result, Err(OAuthError::UserCancelled))); -} - -#[tokio::test] -async fn eliciting_oauth_handler_emits_completion_after_callback() { - let (event_tx, mut event_rx) = mpsc::channel(4); - let handler = Arc::new( - ElicitingOAuthHandler::new(OAuthHandlerContext { - server_name: "linear".to_string(), - callback_port: None, - tx: event_tx, - }) - .unwrap(), - ); - let callback_url = handler.redirect_uri().to_string(); - let authorize = { - let handler = handler.clone(); - tokio::spawn(async move { handler.authorize("https://linear.app/oauth").await }) - }; - - let McpClientEvent::Elicitation(request) = event_rx.recv().await.unwrap() else { - panic!("expected OAuth elicitation"); - }; - request.response_sender.send(rmcp::model::ElicitResult::new(rmcp::model::ElicitationAction::Accept)).unwrap(); - let port = callback_url - .strip_prefix("http://localhost:") - .and_then(|value| value.strip_suffix('/')) - .unwrap() - .parse::() - .unwrap(); - let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap(); - stream.write_all(b"GET /?code=test-code&state=test-state HTTP/1.1\r\nHost: localhost\r\n\r\n").await.unwrap(); - authorize.await.unwrap().unwrap(); - - let completion = event_rx.try_recv().expect("OAuth callback should complete the URL elicitation"); - assert!(matches!( - completion, - McpClientEvent::ElicitationComplete { server_name, elicitation_id } - if server_name == "linear" && elicitation_id == "aether-oauth" - )); -} - -#[tokio::test] -async fn builder_with_oauth_handler_factory_spawns_successfully() { - let handler = Arc::new(CancellingOAuthHandler); - - let mut spawn = mcp("/workspace") - .with_oauth_handler_factory(Arc::new(move |_ctx| Ok(handler.clone()))) - .with_servers(vec![]) - .spawn() - .await - .unwrap(); - let snapshot = spawn.block_until_ready().await.expect("bootstrap completes"); - - assert_eq!(snapshot.tool_definitions().len(), 0); - assert!(snapshot.model_instructions().is_empty()); -} - -#[tokio::test] -async fn http_server_without_handler_stashes_failed_status() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(false); - assert!(manager.add_mcps(vec![endpoint.server("test_server", ToolExposure::ModelVisible)]).await.is_ok()); -} - -#[tokio::test] -async fn real_401_challenge_is_classified_as_needs_oauth() { - let endpoint = UnauthorizedHttpEndpoint::bind( - r#"Bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource", scope="tools:read""#, - ) - .await; - let mut manager = test_manager(true); - - manager.add_mcps(vec![endpoint.server("protected")]).await.unwrap(); - - let status = &manager.server_statuses()[0]; - assert!(matches!(status.status, McpServerStatus::NeedsOAuth)); - assert_eq!(status.auth_capability, McpServerAuthCapability::OAuth); - assert!(status.can_authenticate()); -} - -#[tokio::test] -async fn http_server_with_handler_classifies_non_auth_failure_as_failed() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(true); - - assert!(manager.add_mcps(vec![endpoint.server("test_oauth_server", ToolExposure::ModelVisible)]).await.is_ok()); - - let statuses = manager.server_statuses(); - assert_eq!(statuses.len(), 1); - assert_eq!(statuses[0].name, "test_oauth_server"); - assert!( - matches!(statuses[0].status, McpServerStatus::Failed { .. }), - "Expected Failed, got: {:?}", - statuses[0].status - ); - assert!(!statuses[0].can_authenticate()); -} - -#[tokio::test] -async fn add_mcps_continues_on_oauth_failure() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(true); - - assert!( - manager - .add_mcps(vec![ - endpoint.server("failing_server_1", ToolExposure::ModelVisible), - endpoint.server("failing_server_2", ToolExposure::ModelVisible) - ]) - .await - .is_ok() - ); - assert_eq!(manager.tool_definitions().len(), 0); -} - -#[tokio::test] -async fn accept_oauth_callback_parses_code_and_state() { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let callback_url = format!("http://127.0.0.1:{port}/oauth2callback?code=abc123&state=csrf_token"); - - let handle = tokio::spawn(async move { accept_oauth_callback(&listener).await }); - - let client = reqwest::Client::new(); - let _response = client.get(&callback_url).send().await.expect("Failed to send callback request"); - - let result = handle.await.unwrap(); - let callback_url = result.unwrap(); - let callback = reqwest::Url::parse(&callback_url).unwrap(); - let params = callback.query_pairs().collect::>(); - assert_eq!(params.get("code").map(AsRef::as_ref), Some("abc123")); - assert_eq!(params.get("state").map(AsRef::as_ref), Some("csrf_token")); -} - -#[tokio::test] -async fn deferred_server_with_failing_http_surfaces_failure() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(true); - - let servers = vec![ - fake_mcp("local", FakeMcpServer::new()).with_exposure(ToolExposure::deferred_all()), - endpoint.server("remote", ToolExposure::deferred_all()), - ]; - - let _ = manager.add_mcps(servers).await; - let statuses = manager.server_statuses(); - - let remote_status = statuses.iter().find(|s| s.name == "remote").expect("Expected status entry for 'remote'"); - assert!( - matches!(remote_status.status, McpServerStatus::Failed { .. }), - "Expected Failed for failing HTTP server, got: {:?}", - remote_status.status - ); - assert_eq!(remote_status.auth_capability, McpServerAuthCapability::Unavailable); - assert!(!remote_status.can_authenticate()); - assert!(remote_status.deferred_tools, "Expected remote to have deferred tools"); - - let local_status = statuses.iter().find(|s| s.name == "local").expect("Expected status entry for 'local'"); - assert!(matches!(local_status.status, McpServerStatus::Connected { .. })); - assert!(local_status.deferred_tools); - assert_eq!(manager.tool_definitions().len(), 0); -} - -#[tokio::test] -async fn selective_policy_survives_reconnection_after_failure() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(true); - let selective = endpoint.server("remote", ToolExposure::Deferred(DeferredToolRules::new(&[], &["add_*"]))); - manager.add_mcps(vec![selective]).await.unwrap(); - - assert!(manager.authenticate_server_task("remote").await.is_err()); - - let connected = fake_mcp("remote", FakeMcpServer::new()).with_exposure(ToolExposure::deferred_all()); - let attempt = manager.connect_pending_task(connected).await; - manager.apply_connection_attempt(attempt).await; - - let remote = manager.server_statuses().into_iter().find(|status| status.name == "remote").unwrap(); - assert!(matches!(remote.status, McpServerStatus::Connected { .. })); - assert!(remote.deferred_tools); - let names = manager.tool_definitions().into_iter().map(|tool| tool.name).collect::>(); - assert_eq!(names, ["remote__add_numbers"]); -} - -#[tokio::test] -async fn deferred_tools_partial_connection_works() { - let endpoint = FailingHttpEndpoint::bind().await; - let mut manager = test_manager(false); - - let servers = vec![ - fake_mcp("working", FakeMcpServer::new()).with_exposure(ToolExposure::deferred_all()), - endpoint.server("broken", ToolExposure::deferred_all()), - ]; - - let _ = manager.add_mcps(servers).await; - - assert_eq!(manager.tool_definitions().len(), 0); - let statuses = manager.server_statuses(); - assert!(matches!( - statuses.iter().find(|status| status.name == "working").unwrap().status, - McpServerStatus::Connected { .. } - )); - assert!(matches!( - statuses.iter().find(|status| status.name == "broken").unwrap().status, - McpServerStatus::Failed { .. } - )); -} diff --git a/crates/aether-core/tests/mcp/server_round_trip_tests.rs b/crates/aether-core/tests/mcp/server_round_trip_tests.rs index 166d840e0..a04546563 100644 --- a/crates/aether-core/tests/mcp/server_round_trip_tests.rs +++ b/crates/aether-core/tests/mcp/server_round_trip_tests.rs @@ -1,4 +1,5 @@ -use aether_core::testing::McpTestBuilder; +use super::support::call_tool; +use aether_core::testing::test_agent; use mcp_servers::ReviewMcp; use rmcp::model::{ElicitResult, ElicitationAction}; use serde_json::json; @@ -12,26 +13,17 @@ async fn artifact_review_round_trips_complete_feedback_through_the_production_ex let context = "# Markdown review feedback\n\n## `plan.md`\n\n### Lines 3-4 — Paragraph\n\n```markdown\nShip it.\n```\n\n> Add tests.\n\n### Heading: Plan\n\n> Keep rollout staged."; for suffix in ["", "\n", "\n\n"] { let feedback = format!("{context}{suffix}"); - let test = McpTestBuilder::new() - .server("review", review_mcp_at(root.path())) - .elicitation_response( - ElicitResult::new(ElicitationAction::Accept) - .with_content(json!({"decision": "feedback", "feedback": feedback})), - ) - .build() - .await; - let result = test - .call( - "review", - "review_artifact", - json!({"format": "markdown", "source": {"type": "file", "path": "plan.md"}}), - ) - .await - .result - .expect("review completes"); + let agent = test_agent().mcp_server("review", review_mcp_at(root.path())).elicitation_response( + ElicitResult::new(ElicitationAction::Accept) + .with_content(json!({"decision": "feedback", "feedback": feedback})), + ); + + let run = call_tool(agent, REVIEW_ARTIFACT, markdown_plan()).await; + + let result = run.result().expect("review completes"); let output: serde_json::Value = noyalib::from_str(&result.result).unwrap(); assert_eq!(output, json!({"status": "feedback", "feedback": feedback}), "encoded result: {:?}", result.result); - assert_eq!(test.elicitations().len(), 1); + assert_eq!(run.elicitations.len(), 1); } } @@ -45,21 +37,11 @@ async fn artifact_review_round_trips_approval_and_cancellation_through_the_produ (ElicitResult::new(ElicitationAction::Cancel), "cancelled"), (ElicitResult::new(ElicitationAction::Decline), "declined"), ] { - let test = McpTestBuilder::new() - .server("review", review_mcp_at(root.path())) - .elicitation_response(response) - .build() - .await; - let result = test - .call( - "review", - "review_artifact", - json!({"format": "markdown", "source": {"type": "file", "path": "plan.md"}}), - ) - .await - .result - .expect("review completes"); - let output: serde_json::Value = noyalib::from_str(&result.result).unwrap(); + let agent = test_agent().mcp_server("review", review_mcp_at(root.path())).elicitation_response(response); + + let run = call_tool(agent, REVIEW_ARTIFACT, markdown_plan()).await; + + let output: serde_json::Value = noyalib::from_str(&run.result().expect("review completes").result).unwrap(); assert_eq!(output, json!({"status": status})); assert_eq!(std::fs::read_to_string(&path).unwrap(), "# Plan"); } @@ -74,33 +56,24 @@ async fn artifact_review_round_trips_html_annotations_through_the_production_exe "annotations": [{"element": "

", "excerpt": "Ship faster", "comment": "Make this larger."}] }); let browser = submission.clone(); - let test = McpTestBuilder::new() - .server("review", review_mcp_at(root.path())) + let agent = test_agent() + .mcp_server("review", review_mcp_at(root.path())) .on_url_elicitation(move |url, token| { let body = browser.clone(); async move { submit_in_browser(&url, &token, body).await } }) - .elicitation_response(ElicitResult::new(ElicitationAction::Accept)) - .build() - .await; + .elicitation_response(ElicitResult::new(ElicitationAction::Accept)); + let page = json!({ + "format": "html", + "title": "Landing page", + "source": {"type": "content", "content": "

Ship faster

"} + }); - let result = test - .call( - "review", - "review_artifact", - json!({ - "format": "html", - "title": "Landing page", - "source": {"type": "content", "content": "

Ship faster

"} - }), - ) - .await - .result - .expect("review completes"); + let run = call_tool(agent, REVIEW_ARTIFACT, page).await; - let output: serde_json::Value = noyalib::from_str(&result.result).unwrap(); + let output: serde_json::Value = noyalib::from_str(&run.result().expect("review completes").result).unwrap(); assert_eq!(output, submission); - assert_eq!(test.elicitations().len(), 1); + assert_eq!(run.elicitations.len(), 1); } async fn submit_in_browser(url: &str, token: &str, body: serde_json::Value) { @@ -115,5 +88,11 @@ async fn submit_in_browser(url: &str, token: &str, body: serde_json::Value) { } fn review_mcp_at(root: &Path) -> ReviewMcp { - ReviewMcp::from_args_with_base_dir(Vec::new(), root).expect("empty review MCP arguments are valid") + ReviewMcp::new(root.to_path_buf()) } + +fn markdown_plan() -> serde_json::Value { + json!({"format": "markdown", "source": {"type": "file", "path": "plan.md"}}) +} + +const REVIEW_ARTIFACT: &str = "review__review_artifact"; diff --git a/crates/aether-core/tests/mcp/support.rs b/crates/aether-core/tests/mcp/support.rs new file mode 100644 index 000000000..a0d9617e3 --- /dev/null +++ b/crates/aether-core/tests/mcp/support.rs @@ -0,0 +1,96 @@ +use aether_core::events::{AgentEvent, ToolEvent}; +use aether_core::mcp::{DeferredToolsClient, McpCatalog}; +use aether_core::testing::{TestAgentBuilder, TestScenario}; +use llm::testing::llm_response; +use llm::{ToolCallError, ToolCallResult}; +use mcp_utils::testing::CapturedElicitation; +use serde_json::Value; + +pub const CALL_ID: &str = "call-1"; + +pub struct ToolRun { + pub events: Vec, + pub elicitations: Vec, +} + +#[derive(Debug, PartialEq)] +pub enum TaskOutcome { + Completed(String), + Failed(String), + Cancelled, +} + +pub async fn deferred_tools(mcp: &McpCatalog) -> Vec { + let Some(socket) = mcp.deferred_tools_socket() else { return Vec::new() }; + let client = DeferredToolsClient::connect(socket.to_path_buf()).await.expect("connect to the deferred tools"); + let servers = client.servers().await.expect("list the deferred servers"); + servers + .iter() + .flat_map(|server| server.tools.iter().map(|tool| format!("{}__{}", server.name, tool.name))) + .collect() +} + +pub async fn call_tool(agent: TestAgentBuilder, tool: &str, arguments: Value) -> ToolRun { + let arguments = arguments.to_string(); + let result = agent + .llm_responses(&[ + llm_response().tool_call(CALL_ID, tool, &[&arguments]).build(), + llm_response().text(&["called"]).build(), + llm_response().text(&["handled the task outcome"]).build(), + ]) + .scenario(TestScenario::new().user_text("call the tool").wait_for(finishes_call)) + .run_with_context() + .await + .expect("the test agent runs"); + ToolRun { events: result.messages, elicitations: result.elicitations } +} + +impl ToolRun { + pub fn result(&self) -> Result<&ToolCallResult, &ToolCallError> { + self.tool_events() + .find_map(|event| match event { + ToolEvent::Result { result, .. } => Some(Ok(result)), + ToolEvent::Error { error } => Some(Err(error)), + _ => None, + }) + .expect("the tool call finished in the foreground") + } + + pub fn created_task(&self) -> Option<&str> { + self.tool_events().find_map(|event| match event { + ToolEvent::TaskCreated { task_id, .. } => Some(task_id.as_str()), + _ => None, + }) + } + + pub fn task_outcome(&self) -> TaskOutcome { + self.tool_events() + .find_map(|event| match event { + ToolEvent::TaskCompleted { result, .. } => Some(TaskOutcome::Completed(result.result.clone())), + ToolEvent::TaskFailed { error, .. } => Some(TaskOutcome::Failed(error.error.clone())), + ToolEvent::TaskCancelled { .. } => Some(TaskOutcome::Cancelled), + _ => None, + }) + .expect("the tool call started a background task that finished") + } + + pub fn tool_events(&self) -> impl Iterator { + self.events.iter().filter_map(|event| match event { + AgentEvent::Tool(event) => Some(event), + _ => None, + }) + } +} + +fn finishes_call(event: &AgentEvent) -> bool { + match event { + AgentEvent::Tool(ToolEvent::Result { result, .. }) => result.id == CALL_ID, + AgentEvent::Tool(ToolEvent::Error { error }) => error.id == CALL_ID, + AgentEvent::Tool( + ToolEvent::TaskCompleted { request, .. } + | ToolEvent::TaskFailed { request, .. } + | ToolEvent::TaskCancelled { request, .. }, + ) => request.id == CALL_ID, + _ => false, + } +} diff --git a/crates/aether-core/tests/mcp/task_tests.rs b/crates/aether-core/tests/mcp/task_tests.rs index 98ec2f66b..473a6d696 100644 --- a/crates/aether-core/tests/mcp/task_tests.rs +++ b/crates/aether-core/tests/mcp/task_tests.rs @@ -1,381 +1,64 @@ -use aether_core::testing::{FakeMcpServer, FakeTool, FakeToolResponse, McpTestBuilder}; -use mcp_utils::client::ToolCallEvent; -use rmcp::model::{ - CallToolResult, CreateTaskResult, DetailedTask, ElicitRequest, ElicitRequestParams, ElicitResult, - ElicitationAction, InputRequest, InputRequests, Task, TaskPayload, TaskStatus, -}; +use super::support::{TaskOutcome, call_tool}; +use aether_core::events::ToolEvent; +use aether_core::testing::test_agent; +use mcp_utils::testing::{FakeMcpServer, FakeTask, completed_task, failed_task, task, working_task}; +use rmcp::model::{DetailedTask, TaskPayload, TaskStatus}; use serde_json::json; use std::time::Duration; -fn task(task_id: &str, status: TaskStatus) -> Task { - let now = chrono::Utc::now().to_rfc3339(); - Task::new(task_id, status, now.clone(), now).with_poll_interval_ms(10) -} - -fn task_server(task_id: &str, states: impl IntoIterator) -> FakeMcpServer { - let seed = task(task_id, TaskStatus::Working); - FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(seed)))) - .with_task(task_id, states) -} - -fn completed_task(task_id: &str, text: &str) -> DetailedTask { - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text(text)]); - DetailedTask::new( - task(task_id, TaskStatus::Completed), - TaskPayload::Completed { result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap() }, - ) -} - -fn input_required_task(task_id: &str, key: &str) -> DetailedTask { - let request = InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::FormElicitationParams { - meta: None, - message: "Provide a name".to_string(), - requested_schema: serde_json::from_value(json!({ - "type": "object", - "properties": { "name": { "type": "string" } }, - "required": ["name"] - })) - .unwrap(), - })); - DetailedTask::new( - task(task_id, TaskStatus::InputRequired), - TaskPayload::InputRequired { input_requests: InputRequests::from([(key.to_string(), request)]) }, - ) -} - #[tokio::test] -async fn task_input_is_elicited_and_sent_to_tasks_update() { - let server = task_server( - "input-task", - [input_required_task("input-task", "answer"), completed_task("input-task", "accepted")], - ); - let state = server.state(); - let response = ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": "Ada" })); - let expected_response = serde_json::to_value(&response).unwrap(); - let test = McpTestBuilder::new().server("tasks", server).elicitation_response(response).build().await; +async fn a_created_task_runs_in_the_background_until_its_outcome_arrives() { + let server = task_server([working_task(TASK), completed_task(TASK, "finished")]); - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task notification"); + let run = call_tool(test_agent().mcp_server("tasks", server), DEFERRED, json!({})).await; - assert_eq!(notification.status, "completed"); - assert!(notification.body.contains("accepted")); - assert_eq!(test.elicitations().len(), 1); - let updates = state.task_updates(); - assert_eq!(updates.len(), 1); - assert_eq!(updates[0].task_id, "input-task"); - assert_eq!(updates[0].input_responses.get("answer"), Some(&expected_response)); + assert_eq!(run.created_task(), Some(TASK)); + assert!(matches!(run.task_outcome(), TaskOutcome::Completed(body) if body.contains("finished"))); } #[tokio::test] -async fn repeated_task_input_keys_terminate_with_an_error() { - let server = task_server( - "repeated-input", - [ - input_required_task("repeated-input", "answer"), - input_required_task("repeated-input", "answer"), - completed_task("repeated-input", "must not complete"), - ], - ); - let state = server.state(); - let test = McpTestBuilder::new() - .server("tasks", server) - .elicitation_response(ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": "Ada" }))) - .build() - .await; +async fn a_failed_task_becomes_a_failed_outcome() { + let server = task_server([failed_task(TASK, "task exploded")]); - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task error notification"); + let run = call_tool(test_agent().mcp_server("tasks", server), DEFERRED, json!({})).await; - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("repeated input requests"), "{}", notification.body); - assert_eq!(test.elicitations().len(), 1); - assert_eq!(state.task_updates().len(), 1); - assert_eq!(state.task_cancel_ids(), ["repeated-input"]); + assert!(matches!(run.task_outcome(), TaskOutcome::Failed(body) if body.contains("task exploded"))); } #[tokio::test] -async fn any_repeated_task_input_key_terminates_without_partial_update() { - let TaskPayload::InputRequired { input_requests: mut second_requests } = - input_required_task("mixed-input", "answer").payload - else { - unreachable!() - }; - let TaskPayload::InputRequired { mut input_requests } = input_required_task("mixed-input", "new-answer").payload - else { - unreachable!() - }; - let new_request = input_requests.remove("new-answer").unwrap(); - second_requests.insert("new-answer".to_string(), new_request); - let second_round = DetailedTask::new( - task("mixed-input", TaskStatus::InputRequired), - TaskPayload::InputRequired { input_requests: second_requests }, - ); - let server = task_server( - "mixed-input", - [ - input_required_task("mixed-input", "answer"), - second_round, - completed_task("mixed-input", "must not complete"), - ], - ); +async fn task_statuses_are_reported_before_the_outcome() { + let halfway = + DetailedTask::new(task(TASK, TaskStatus::Working).with_status_message("halfway"), TaskPayload::Working); + let server = task_server([halfway, completed_task(TASK, "done")]); + + let run = call_tool(test_agent().mcp_server("tasks", server), DEFERRED, json!({})).await; + + let status = run.tool_events().position(|event| { + matches!( + event, + ToolEvent::TaskStatus { task_id, status, status_message, .. } + if task_id == TASK && status == "working" && status_message.as_deref() == Some("halfway") + ) + }); + let outcome = run.tool_events().position(|event| matches!(event, ToolEvent::TaskCompleted { .. })); + assert!(status.expect("task status event") < outcome.expect("task outcome")); +} + +#[tokio::test(start_paused = true)] +async fn the_tool_timeout_bounds_background_tasks() { + let server = task_server([working_task(TASK)]); let state = server.state(); - let test = McpTestBuilder::new() - .server("tasks", server) - .elicitation_response(ElicitResult::new(ElicitationAction::Accept)) - .build() - .await; + let agent = test_agent().mcp_server("tasks", server).tool_timeout(Duration::from_millis(25)); - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task error notification"); + let run = call_tool(agent, DEFERRED, json!({})).await; - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("repeated input requests"), "{}", notification.body); - assert_eq!(test.elicitations().len(), 1); - assert_eq!(state.task_updates().len(), 1, "must not submit a partial second input round"); - assert_eq!(state.task_cancel_ids(), ["mixed-input"]); + assert!(matches!(run.task_outcome(), TaskOutcome::Failed(body) if body.contains("timeout"))); + assert_eq!(state.task_cancel_ids(), [TASK]); } -#[tokio::test] -async fn failed_task_payload_produces_error_notification() { - let failed = DetailedTask::new( - task("failed-task", TaskStatus::Failed), - TaskPayload::Failed { - error: serde_json::from_value(json!({ "code": "boom", "message": "task exploded" })).unwrap(), - }, - ); - let test = McpTestBuilder::new().server("tasks", task_server("failed-task", [failed])).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task error notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("task exploded")); - assert_eq!(notification.status, "failed"); +fn task_server(states: impl IntoIterator) -> FakeMcpServer { + FakeTask::new(TASK).states(states).server("deferred") } -#[tokio::test] -async fn cancelled_task_payload_produces_error_notification() { - let cancelled = DetailedTask::new(task("cancelled-task", TaskStatus::Cancelled), TaskPayload::Cancelled); - let test = McpTestBuilder::new().server("tasks", task_server("cancelled-task", [cancelled])).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task cancellation notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("was cancelled")); - assert_eq!(notification.status, "failed"); -} - -#[tokio::test] -async fn malformed_task_result_produces_error_notification() { - let malformed = DetailedTask::new( - task("malformed-task", TaskStatus::Completed), - TaskPayload::Completed { result: serde_json::from_value(json!({ "unexpected": true })).unwrap() }, - ); - let test = McpTestBuilder::new().server("tasks", task_server("malformed-task", [malformed])).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("malformed task notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("malformed result"), "{}", notification.body); -} - -#[tokio::test] -async fn expired_task_terminates_before_polling() { - let mut seed = task("expired-task", TaskStatus::Working); - seed.created_at = "2020-01-01T00:00:00Z".to_string(); - seed.ttl_ms = Some(1); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(seed)))) - .with_task("expired-task", [completed_task("expired-task", "too late")]); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("expired task notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("expired before completion")); - assert_eq!(state.task_get_ids().len(), 0); -} - -#[tokio::test] -async fn task_execution_deadline_terminates_polling() { - let working = DetailedTask::new(task("deadline-task", TaskStatus::Working), TaskPayload::Working); - let server = task_server("deadline-task", [working]); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).tool_timeout(Duration::from_millis(25)).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("deadline notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("execution deadline"), "{}", notification.body); - assert_eq!(state.task_cancel_ids(), ["deadline-task"]); -} - -#[tokio::test] -async fn tasks_get_failure_terminates_task() { - let server = - task_server("get-failure", [completed_task("get-failure", "must not finish")]).with_task_get_failures(1); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task error notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("failed to get task"), "{}", notification.body); - assert_eq!(state.task_get_ids(), ["get-failure"]); - assert_eq!(state.task_cancel_ids(), ["get-failure"]); -} - -#[tokio::test] -async fn tasks_update_failure_terminates_task() { - let server = - task_server("update-failure", [input_required_task("update-failure", "answer")]).with_task_update_failures(1); - let state = server.state(); - let test = McpTestBuilder::new() - .server("tasks", server) - .elicitation_response(ElicitResult::new(ElicitationAction::Accept)) - .build() - .await; - - test.call("tasks", "deferred", json!({})).await; - let notification = test.next_task_outcome().await.expect("task error notification"); - - assert_eq!(notification.status, "failed"); - assert!(notification.body.contains("failed to update task"), "{}", notification.body); - assert_eq!(state.task_updates().len(), 1); - assert_eq!(state.task_cancel_ids(), ["update-failure"]); -} - -#[tokio::test] -async fn intermediate_task_status_is_emitted_before_outcome() { - let server = task_server( - "status-task", - [ - DetailedTask::new( - task("status-task", TaskStatus::Working).with_status_message("halfway"), - TaskPayload::Working, - ), - completed_task("status-task", "done"), - ], - ); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - test.call("tasks", "deferred", json!({})).await; - let event = test.next_tool_event().await.expect("task status event"); - - assert!(matches!( - event, - ToolCallEvent::TaskStatus(task) - if task.task_id == "status-task" - && task.status == TaskStatus::Working - && task.status_message.as_deref() == Some("halfway") - )); -} - -#[tokio::test] -async fn deferred_task_is_polled_and_injected_as_a_success_notification() { - let seed = task("task-1", TaskStatus::Working); - let completed = DetailedTask::new( - task("task-1", TaskStatus::Completed), - TaskPayload::Completed { - result: serde_json::from_value(json!({ - "content": [{"type": "text", "text": "finished"}], - "isError": false - })) - .unwrap(), - }, - ); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(seed.clone())))) - .with_task("task-1", [DetailedTask::new(seed, TaskPayload::Working), completed]); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - let outcome = test.call("tasks", "deferred", json!({})).await; - assert_eq!(outcome.deferred_task.as_ref().map(|task| task.task.task_id.as_str()), Some("task-1")); - assert!(outcome.result.unwrap().result.contains("task-1")); - - let notification = test.next_task_outcome().await.expect("task notification"); - assert_eq!(notification.task_id, "task-1"); - assert_eq!(notification.status, "completed"); - assert!(notification.body.contains("finished")); - assert_eq!(state.task_get_ids(), ["task-1", "task-1"]); -} - -#[tokio::test] -async fn cancelling_deferred_task_stops_polling_and_notifies_server() { - let working = task("task-cancel", TaskStatus::Working); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(working.clone())))) - .with_task("task-cancel", [DetailedTask::new(working, TaskPayload::Working)]); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - let outcome = test.call("tasks", "deferred", json!({})).await; - let tool_id = outcome.result.expect("deferred acknowledgement").id; - test.cancel_tool(&tool_id); - - loop { - let event = test.next_tool_event().await.expect("tool event before cancellation"); - if let ToolCallEvent::Cancelled { task_id } = event { - assert_eq!(task_id.as_deref(), Some("task-cancel")); - break; - } - } - assert_eq!(state.task_cancel_ids(), ["task-cancel"]); -} - -#[tokio::test] -async fn deferred_task_keeps_its_ordered_stream_until_terminal_outcome() { - let seed = task("ordered-terminal-task", TaskStatus::Completed); - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text("already done")]); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("terminal").responds(FakeToolResponse::task(CreateTaskResult::new(seed.clone())))) - .with_task( - "ordered-terminal-task", - [DetailedTask::new( - seed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap(), - }, - )], - ); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - let outcome = test.call("tasks", "terminal", json!({})).await; - assert!(outcome.deferred_task.is_some()); - - assert!(test.next_tool_event().await.is_some(), "deferred tool stream ended before its terminal outcome"); -} - -#[tokio::test] -async fn terminal_seed_is_refreshed_from_tasks_get_before_notification() { - let seed = task("terminal-task", TaskStatus::Completed); - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text("already done")]); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("terminal").responds(FakeToolResponse::task(CreateTaskResult::new(seed.clone())))) - .with_task( - "terminal-task", - [DetailedTask::new( - seed, - TaskPayload::Completed { - result: serde_json::from_value(serde_json::to_value(result).unwrap()).unwrap(), - }, - )], - ); - let state = server.state(); - let test = McpTestBuilder::new().server("tasks", server).build().await; - - let outcome = test.call("tasks", "terminal", json!({})).await; - assert!(outcome.deferred_task.is_some()); - let notification = test.next_task_outcome().await.expect("terminal task notification"); - assert_eq!(notification.status, "completed"); - assert!(notification.body.contains("already done")); - assert_eq!(state.task_get_ids(), ["terminal-task"]); -} +const TASK: &str = "task-1"; +const DEFERRED: &str = "tasks__deferred"; diff --git a/crates/aether-core/tests/mcp/tool_catalog_tests.rs b/crates/aether-core/tests/mcp/tool_catalog_tests.rs new file mode 100644 index 000000000..ece3eaccc --- /dev/null +++ b/crates/aether-core/tests/mcp/tool_catalog_tests.rs @@ -0,0 +1,212 @@ +use super::support::deferred_tools; +use aether_core::core::McpInstructions; +use aether_core::mcp::{McpCatalog, ToolAnnotationMatcher, ToolExposure, ToolFilter, ToolMatcher, mcp}; +use aether_core::testing::ready_catalog; +use futures::StreamExt; +use llm::ToolDefinition; +use mcp_utils::client::{ToolCallOptions, Transport}; +use mcp_utils::gateway::ServerSpec; +use mcp_utils::model::ToolAnnotations; +use mcp_utils::testing::{FakeAuthServer, FakeMcpServer, FakeTool, OAuthPrompt, args, elicitation_sink}; +use serde_json::{Map, json}; +use std::collections::HashMap; +use utils::mcp_status::{McpServerAuthCapability, McpServerStatus}; + +#[tokio::test] +async fn model_tools_are_namespaced_and_keep_their_metadata() { + let server = FakeMcpServer::math() + .with_tool(FakeTool::new("lookup").description("Looks up").annotations(ToolAnnotations::new().read_only(true))); + let mcp = ready_catalog(mcp("/workspace").with_servers(vec![ServerSpec::in_process("math", server)])).await; + + let tools = mcp.tool_definitions(); + + assert_eq!(names(&tools), ["math__add_numbers", "math__divide_numbers", "math__lookup"]); + let lookup = &tools[2]; + assert_eq!(lookup.description, "Looks up"); + assert_eq!(lookup.server.as_deref(), Some("math")); + assert_eq!(lookup.parameters, json!({"type": "object", "properties": {}})); + assert_eq!(lookup.annotations.as_ref().and_then(|annotations| annotations.read_only_hint), Some(true)); +} + +#[tokio::test] +async fn deferred_tools_are_listed_apart_from_model_tools_and_not_callable_by_the_model() { + let mcp = ready_catalog(mcp("/workspace").with_server(math("math"), deferring(&["divide_numbers"]))).await; + + let deferred_as_model = mcp.call_tool("math__divide_numbers", Map::new(), ToolCallOptions::default()); + + assert_eq!(names(&mcp.tool_definitions()), ["math__add_numbers"]); + assert_eq!(deferred_tools(&mcp).await, ["math__divide_numbers"]); + assert!(deferred_as_model.is_none()); + assert!(mcp.statuses()[0].deferred_tools); +} + +#[tokio::test] +async fn calls_reach_servers_whose_names_contain_the_delimiter() { + let (math, math_v2) = (FakeMcpServer::math(), FakeMcpServer::math()); + let (math_state, math_v2_state) = (math.state(), math_v2.state()); + let servers = vec![ServerSpec::in_process("math", math), ServerSpec::in_process("math__v2", math_v2)]; + let mcp = ready_catalog(mcp("/workspace").with_servers(servers)).await; + + let sum = mcp.call_tool("math__v2__add_numbers", args(json!({"a": 2, "b": 3})), ToolCallOptions::default()); + + assert_eq!(sum.unwrap().result().await.unwrap().structured_content, Some(json!({"sum": 5}))); + assert!(math_state.calls_for("add_numbers").is_empty()); + assert_eq!(math_v2_state.calls_for("add_numbers").len(), 1); +} + +#[tokio::test] +async fn unknown_tools_are_not_offered() { + let mcp = ready_catalog(mcp("/workspace").with_servers(vec![math("math")])).await; + + for name in ["math__missing", "other__add_numbers", "mathadd_numbers", "add_numbers"] { + assert!(mcp.call_tool(name, Map::new(), ToolCallOptions::default()).is_none(), "{name} was offered"); + } +} + +#[tokio::test] +async fn the_tool_filter_hides_model_and_deferred_tools() { + let server = + FakeMcpServer::math().with_tool(FakeTool::new("peek").annotations(ToolAnnotations::new().read_only(true))); + let filter = ToolFilter { + allow: Vec::new(), + deny: vec![ + ToolMatcher::name("direct__divide_*"), + ToolMatcher::annotations(ToolAnnotationMatcher { + read_only: Some(true), + ..ToolAnnotationMatcher::default() + }), + ], + }; + let mcp = ready_catalog( + mcp("/workspace") + .with_servers(vec![ServerSpec::in_process("direct", server.clone())]) + .with_server(ServerSpec::in_process("lazy", server), ToolExposure::deferred_all()) + .with_tool_filter(filter), + ) + .await; + + assert_eq!(names(&mcp.tool_definitions()), ["direct__add_numbers"]); + assert_eq!(deferred_tools(&mcp).await, ["lazy__add_numbers", "lazy__divide_numbers"]); + assert!(mcp.call_tool("direct__peek", Map::new(), ToolCallOptions::default()).is_none()); + let statuses = mcp.statuses(); + assert_eq!(statuses[0].status, McpServerStatus::Connected { tool_count: 3 }); + assert!(!statuses[0].deferred_tools && statuses[1].deferred_tools); +} + +#[tokio::test] +async fn instructions_come_only_from_servers_offering_model_tools() { + let server = FakeMcpServer::math().with_instructions("Use the math tools"); + let mcp = ready_catalog( + mcp("/workspace") + .with_servers(vec![ServerSpec::in_process("direct", server.clone())]) + .with_server(ServerSpec::in_process("lazy", server), ToolExposure::deferred_all()), + ) + .await; + + let instructions = mcp.instructions(); + + assert_eq!(instructions.servers, [("direct".to_string(), "Use the math tools".to_string())].into()); + assert!(instructions.progressive_discovery); +} + +#[tokio::test] +async fn servers_whose_tools_are_all_filtered_out_give_no_instructions() { + let server = FakeMcpServer::math().with_instructions("Use the math tools"); + let filter = ToolFilter { allow: vec![ToolMatcher::name("kept__*")], deny: Vec::new() }; + let servers = vec![ServerSpec::in_process("kept", server.clone()), ServerSpec::in_process("hidden", server)]; + let mcp = ready_catalog(mcp("/workspace").with_servers(servers).with_tool_filter(filter)).await; + + let instructions = mcp.instructions(); + + assert_eq!(instructions.servers.keys().collect::>(), ["kept"]); +} + +#[tokio::test] +async fn deferred_servers_that_are_not_connected_need_no_discovery_instructions() { + let mcp = ready_catalog(mcp("/workspace").with_server(missing_binary("lazy"), ToolExposure::deferred_all())).await; + + assert!(!mcp.instructions().progressive_discovery); +} + +#[tokio::test] +async fn statuses_report_each_server_with_its_exposure_and_sign_in_support() { + let auth = FakeAuthServer::start(FakeMcpServer::math()).await; + let (sink, _host) = elicitation_sink(); + let mcp = ready_catalog( + mcp("/workspace") + .with_servers(vec![math("math"), auth.spec("linear")]) + .with_server(missing_binary("broken"), ToolExposure::deferred_all()) + .with_elicitation(sink), + ) + .await; + + let statuses = mcp.statuses(); + + let status = |name: &str| statuses.iter().find(|entry| entry.name == name).expect("a status for every server"); + assert_eq!(status("math").status, McpServerStatus::Connected { tool_count: 2 }); + assert_eq!(status("math").auth_capability, McpServerAuthCapability::Unavailable); + assert!(!status("math").deferred_tools); + assert_eq!(status("linear").status, McpServerStatus::NeedsOAuth); + assert_eq!(status("linear").auth_capability, McpServerAuthCapability::OAuth); + assert!(matches!(&status("broken").status, McpServerStatus::Failed { error } if error.contains("Failed to spawn"))); + assert_eq!(status("broken").auth_capability, McpServerAuthCapability::Unavailable); + assert!(status("broken").deferred_tools); +} + +#[tokio::test] +async fn signing_in_connects_a_server_with_its_configured_exposure() { + let auth = FakeAuthServer::start(FakeMcpServer::math()).await; + let (sink, mut host) = elicitation_sink(); + let mcp = ready_catalog( + mcp("/workspace").with_server(auth.spec("linear"), deferring(&["divide_numbers"])).with_elicitation(sink), + ) + .await; + let signing_in = tokio::spawn({ + let mcp = mcp.clone(); + async move { mcp.authenticate("linear").await } + }); + + let prompt = OAuthPrompt::receive(&mut host).await; + assert_eq!(mcp.statuses()[0].status, McpServerStatus::Authenticating); + prompt.approve().await; + signing_in.await.unwrap().expect("signing in succeeds"); + + assert_eq!(mcp.statuses()[0].status, McpServerStatus::Connected { tool_count: 2 }); + assert_eq!(names(&mcp.tool_definitions()), ["linear__add_numbers"]); + assert_eq!(deferred_tools(&mcp).await, ["linear__divide_numbers"]); +} + +#[tokio::test] +async fn the_default_catalog_offers_nothing_and_never_changes() { + let mcp = McpCatalog::default(); + + assert!(mcp.tool_definitions().is_empty() && mcp.statuses().is_empty()); + assert_eq!(mcp.instructions(), McpInstructions::default()); + assert!(mcp.deferred_tools_socket().is_none()); + assert!(mcp.changes().next().await.is_none()); +} + +fn math(name: &str) -> ServerSpec { + ServerSpec::in_process(name, FakeMcpServer::math()) +} + +fn missing_binary(name: &str) -> ServerSpec { + let transport = Transport::Stdio { + command: "aether-missing-mcp-server".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, + }; + ServerSpec { name: name.to_string(), transport } +} + +fn deferring(tools: &[&str]) -> ToolExposure { + ToolExposure::Deferred(ToolFilter { + allow: tools.iter().copied().map(ToolMatcher::name).collect(), + deny: Vec::new(), + }) +} + +fn names(tools: &[ToolDefinition]) -> Vec<&str> { + tools.iter().map(|tool| tool.name.as_str()).collect() +} diff --git a/crates/aether-core/tests/mcp/tool_filter_tests.rs b/crates/aether-core/tests/mcp/tool_filter_tests.rs new file mode 100644 index 000000000..c1430fb99 --- /dev/null +++ b/crates/aether-core/tests/mcp/tool_filter_tests.rs @@ -0,0 +1,140 @@ +use aether_core::mcp::{ToolAnnotationMatcher, ToolFilter, ToolMatcher}; +use mcp_utils::model::{Tool, ToolAnnotations}; +use std::sync::Arc; + +#[test] +fn empty_filter_allows_all_tools() { + assert_eq!(allowed(&ToolFilter::default(), [tool("bash"), tool("read_file")]), ["bash", "read_file"]); +} + +#[test] +fn allow_keeps_only_matching_tools() { + let filter = ToolFilter { allow: vec![ToolMatcher::name("read_file"), ToolMatcher::name("grep")], deny: vec![] }; + assert_eq!(allowed(&filter, [tool("bash"), tool("read_file"), tool("grep")]), ["read_file", "grep"]); +} + +#[test] +fn deny_removes_matching_tools() { + let filter = ToolFilter { allow: vec![], deny: vec![ToolMatcher::name("bash")] }; + assert_eq!(allowed(&filter, [tool("bash"), tool("read_file")]), ["read_file"]); +} + +#[test] +fn wildcard_matching() { + let filter = ToolFilter { allow: vec![ToolMatcher::name("coding__*")], deny: vec![] }; + let tools = [tool("coding__grep"), tool("coding__read_file"), tool("plugins__bash")]; + assert_eq!(allowed(&filter, tools), ["coding__grep", "coding__read_file"]); +} + +#[test] +fn combined_allow_and_deny() { + let filter = + ToolFilter { allow: vec![ToolMatcher::name("coding__*")], deny: vec![ToolMatcher::name("coding__write_file")] }; + let tools = [tool("coding__grep"), tool("coding__write_file"), tool("coding__read_file"), tool("plugins__bash")]; + assert_eq!(allowed(&filter, tools), ["coding__grep", "coding__read_file"]); +} + +#[test] +fn annotation_allow_matches_present_values() { + let filter = ToolFilter { allow: vec![ToolMatcher::read_only()], deny: vec![] }; + let tools = [ + tool("unknown"), + annotated("read", ToolAnnotations::new().read_only(true)), + annotated("write", ToolAnnotations::new().read_only(false)), + ]; + assert_eq!(allowed(&filter, tools), ["read"]); +} + +#[test] +fn deny_annotation_removes_destructive_tools() { + let filter = ToolFilter { + allow: vec![], + deny: vec![annotation(ToolAnnotationMatcher { destructive: Some(true), ..ToolAnnotationMatcher::default() })], + }; + let tools = [ + tool("unknown"), + annotated("safe_update", ToolAnnotations::new().read_only(false).destructive(false)), + annotated("delete", ToolAnnotations::new().destructive(true)), + ]; + assert_eq!(allowed(&filter, tools), ["unknown", "safe_update"]); +} + +#[test] +fn annotation_matchers_do_not_match_missing_fields() { + let filter = ToolFilter { + allow: vec![], + deny: vec![ + annotation(ToolAnnotationMatcher { destructive: Some(true), ..ToolAnnotationMatcher::default() }), + annotation(ToolAnnotationMatcher { open_world: Some(true), ..ToolAnnotationMatcher::default() }), + annotation(ToolAnnotationMatcher { idempotent: Some(false), ..ToolAnnotationMatcher::default() }), + annotation(ToolAnnotationMatcher { read_only: Some(false), ..ToolAnnotationMatcher::default() }), + ], + }; + assert_eq!(allowed(&filter, [tool("unknown")]), ["unknown"]); +} + +#[test] +fn annotation_matchers_do_not_infer_fields_from_read_only_hint() { + let filter = ToolFilter { + allow: vec![annotation(ToolAnnotationMatcher { destructive: Some(false), ..ToolAnnotationMatcher::default() })], + deny: vec![], + }; + assert_eq!(allowed(&filter, [annotated("read", ToolAnnotations::new().read_only(true))]).len(), 0); +} + +#[test] +fn deny_wins_over_allow() { + let filter = + ToolFilter { allow: vec![ToolMatcher::read_only()], deny: vec![ToolMatcher::name("coding__read_file")] }; + assert_eq!(allowed(&filter, [annotated("coding__read_file", ToolAnnotations::new().read_only(true))]).len(), 0); +} + +#[test] +fn mixed_allow_entries_are_ored() { + let filter = ToolFilter { allow: vec![ToolMatcher::read_only(), ToolMatcher::name("review__*")], deny: vec![] }; + let tools = [ + annotated("coding__grep", ToolAnnotations::new().read_only(true)), + tool("review__review_artifact"), + tool("coding__bash"), + ]; + assert_eq!(allowed(&filter, tools), ["coding__grep", "review__review_artifact"]); +} + +#[test] +fn empty_annotation_matcher_matches_nothing() { + let filter = ToolFilter { allow: vec![annotation(ToolAnnotationMatcher::default())], deny: vec![] }; + assert_eq!(allowed(&filter, [annotated("coding__grep", ToolAnnotations::new().read_only(true))]).len(), 0); +} + +#[test] +fn exact_name_match_is_not_a_prefix_match() { + let filter = ToolFilter { allow: vec![ToolMatcher::name("bash")], deny: vec![] }; + assert_eq!(allowed(&filter, [tool("bash"), tool("bash_extended")]), ["bash"]); +} + +#[test] +fn tool_matcher_uses_exact_and_trailing_wildcard_names() { + let exact = ToolMatcher::name("foo"); + let wildcard = ToolMatcher::name("foo*"); + assert!(exact.matches("foo", &tool("foo"))); + assert!(!exact.matches("foobar", &tool("foobar"))); + assert!(wildcard.matches("foobar", &tool("foobar"))); + assert!(wildcard.matches("foo", &tool("foo"))); + assert!(!wildcard.matches("bar", &tool("bar"))); +} + +fn allowed(filter: &ToolFilter, tools: impl IntoIterator) -> Vec { + tools.into_iter().filter(|tool| filter.allows(&tool.name, tool)).map(|tool| tool.name.to_string()).collect() +} + +fn tool(name: &str) -> Tool { + Tool::new(name.to_string(), "", Arc::new(serde_json::Map::new())) +} + +fn annotated(name: &str, annotations: ToolAnnotations) -> Tool { + tool(name).with_annotations(annotations) +} + +fn annotation(matcher: ToolAnnotationMatcher) -> ToolMatcher { + ToolMatcher::annotations(matcher) +} diff --git a/crates/aether-core/tests/mcp/trace_context_tests.rs b/crates/aether-core/tests/mcp/trace_context_tests.rs index 88b98598c..d753031d6 100644 --- a/crates/aether-core/tests/mcp/trace_context_tests.rs +++ b/crates/aether-core/tests/mcp/trace_context_tests.rs @@ -1,5 +1,7 @@ -use aether_core::events::TraceContext; -use aether_core::testing::{FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, McpTestBuilder}; +use super::support::call_tool; +use aether_core::events::{AgentEvent, AgentObserver, TraceContext}; +use aether_core::testing::test_agent; +use mcp_utils::testing::{FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse}; use serde_json::json; #[tokio::test] @@ -9,9 +11,9 @@ async fn tool_call_propagates_w3c_trace_context_in_request_metadata() { tracestate: Some("vendor=value".to_string()), }; let (server, state) = capture_server(); - let test = McpTestBuilder::new().server("trace", server).trace_context(trace_context.clone()).build().await; + let agent = test_agent().mcp_server("trace", server).observer(Box::new(Tracing(trace_context.clone()))); - test.call("trace", "capture", json!({})).await.result.expect("capture tool succeeds"); + call_tool(agent, "trace__capture", json!({})).await.result().expect("capture tool succeeds"); let calls = state.calls_for("capture"); assert_eq!(calls[0].context_meta["traceparent"], json!(trace_context.traceparent)); @@ -21,15 +23,27 @@ async fn tool_call_propagates_w3c_trace_context_in_request_metadata() { #[tokio::test] async fn tool_call_without_trace_context_sends_no_trace_metadata() { let (server, state) = capture_server(); - let test = McpTestBuilder::new().server("trace", server).build().await; - test.call("trace", "capture", json!({})).await.result.expect("capture tool succeeds"); + call_tool(test_agent().mcp_server("trace", server), "trace__capture", json!({})) + .await + .result() + .expect("capture tool succeeds"); let calls = state.calls_for("capture"); assert_eq!(calls[0].context_meta.get("traceparent"), None); assert_eq!(calls[0].context_meta.get("tracestate"), None); } +struct Tracing(TraceContext); + +impl AgentObserver for Tracing { + fn on_event(&mut self, _message: &AgentEvent) {} + + fn tool_trace_context(&self, _tool_id: &str) -> Option { + Some(self.0.clone()) + } +} + fn capture_server() -> (FakeMcpServer, FakeMcpState) { let server = FakeMcpServer::new().with_tool(FakeTool::new("capture").responds(FakeToolResponse::text("captured"))); let state = server.state(); diff --git a/crates/aether-project/Cargo.toml b/crates/aether-project/Cargo.toml index 8e4a0e005..f781ad23f 100644 --- a/crates/aether-project/Cargo.toml +++ b/crates/aether-project/Cargo.toml @@ -22,7 +22,6 @@ workspace = true [dependencies] aether-core = { package = "aether-agent-core", path = "../aether-core", version = "0.9.4" } llm = { package = "aether-llm", path = "../llm", version = "0.9.13" } -mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", version = "0.7.13" } utils = { package = "aether-utils", path = "../utils", version = "0.2.19" } schemars = { workspace = true } serde = { workspace = true } diff --git a/crates/aether-project/src/aether_settings.rs b/crates/aether-project/src/aether_settings.rs index 64acebc8f..8ad210bd3 100644 --- a/crates/aether-project/src/aether_settings.rs +++ b/crates/aether-project/src/aether_settings.rs @@ -2,10 +2,9 @@ use utils::SettingsStore; use crate::agent_config::AgentConfig; use crate::error::SettingsError; -use crate::{McpFileSpec, McpSourceSpec, PromptSource}; +use crate::{McpSourceSpec, PromptSource}; use aether_core::core::Prompt; use llm::ProviderConnectionOverrides; -use mcp_utils::client::McpConfig; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; @@ -425,45 +424,11 @@ fn inline_prompt_sources(sources: &[PromptSource], root: &Path) -> Result Result, SettingsError> { - let mut inlined = Vec::new(); - for source in sources { - let McpSourceSpec::File(McpFileSpec { path, defer_tools, optional }) = source else { - inlined.push(source.clone()); - continue; - }; - - let full_path = match path.resolve(root) { - Ok(full_path) => full_path, - Err(VarError::NotFound(variable)) => { - if *optional { - tracing::warn!( - "Skipping optional MCP config '{}': variable '{variable}' is not defined", - path.as_authored() - ); - continue; - } - return Err(SettingsError::UnresolvedMcpConfigVariable { - path: path.as_authored().to_string(), - variable, - }); - } - }; - - if !full_path.is_file() { - if *optional { - continue; - } - return Err(SettingsError::InvalidMcpConfigPath { path: path.as_authored().to_string() }); - } - - let mut config = McpConfig::from_json_file(&full_path) - .map_err(|e| SettingsError::IoError(format!("Failed to read MCP config '{}': {e}", full_path.display())))?; - if *defer_tools { - config.defer_all_tools(); - } - inlined.push(McpSourceSpec::Inline { servers: config.servers }); - } - Ok(inlined) + sources + .iter() + .filter_map(|spec| spec.load(root).transpose()) + .map(|config| config.map(McpSourceSpec::Inline)) + .collect() } fn promote_prompt_sources(sources: &mut [PromptSource], source_root: &Path) { @@ -496,10 +461,10 @@ impl TryFrom<&str> for AetherSettings { #[cfg(test)] mod tests { use super::*; - use crate::testing::{agent_config, agent_json, agent_json_with, home, project, settings_agent}; + use crate::testing::{agent_config, agent_json, agent_json_with, home, mcp_json, project, settings_agent}; use crate::{AgentCatalog, McpFileSpec, McpSourceSpec, PromptSource}; - use aether_core::agent_spec::McpConfigSource; use aether_core::core::Prompt; + use aether_core::mcp::McpConfig; use serde_json::json; use std::collections::BTreeMap; @@ -839,7 +804,7 @@ mod tests { let project = project(); let home = home() .file(".aether/agents/user.md", "User instructions") - .file(".aether/mcp/user.json", r#"{"servers":{}}"#) + .file(".aether/mcp/user.json", &mcp_json("user")) .settings( &json!({ "agents": [agent_json_with( @@ -861,10 +826,7 @@ mod tests { Prompt::File { path, .. } => path == &expected_prompt, Prompt::Text(_) | Prompt::McpInstructions(_) => false, })); - assert!(matches!( - &spec.mcp_config_sources[0], - McpConfigSource::File { path, defer_tools: false } if path == &aether_home.join("mcp/user.json") - )); + assert_eq!(spec.mcp.servers.keys().collect::>(), ["user"]); } #[test] @@ -951,9 +913,9 @@ mod tests { assert_eq!(settings.prompts, vec![PromptSource::Text { text: "Be helpful".to_string() }]); assert_eq!(settings.agents[0].prompts, vec![PromptSource::Text { text: "Edit carefully".to_string() }]); - assert!(matches!(&settings.mcps[0], McpSourceSpec::Inline { servers } if servers.contains_key("coding"))); + assert!(matches!(&settings.mcps[0], McpSourceSpec::Inline(config) if config.servers.contains_key("coding"))); assert!( - matches!(&settings.agents[0].mcps[0], McpSourceSpec::Inline { servers } if servers.contains_key("coding")) + matches!(&settings.agents[0].mcps[0], McpSourceSpec::Inline(config) if config.servers.contains_key("coding")) ); let serialized = serde_json::to_string(&settings).unwrap(); @@ -995,7 +957,7 @@ mod tests { let config = AetherSettings { agent: None, agents: vec![AgentConfig { - mcps: vec![McpSourceSpec::Inline { servers: BTreeMap::new() }], + mcps: vec![McpSourceSpec::Inline(McpConfig::default())], ..agent_config("alpha") }], ..AetherSettings::default() @@ -1004,8 +966,7 @@ mod tests { let catalog = AgentCatalog::from_settings(dir.root(), config).unwrap(); let spec = catalog.resolve("alpha").unwrap(); - assert_eq!(spec.mcp_config_sources.len(), 1); - assert!(matches!(spec.mcp_config_sources[0], McpConfigSource::Inline(_))); + assert!(spec.mcp.is_empty()); } #[test] @@ -1070,14 +1031,6 @@ mod tests { assert_eq!(value["agents"][0]["mcps"], serde_json::json!(["agent-mcp.json"])); } - #[test] - fn deserializes_legacy_proxy_mcp_file_as_deferred() { - let source: McpSourceSpec = - serde_json::from_value(serde_json::json!({"type":"file", "path":"mcp.json", "proxy":true})).unwrap(); - - assert!(matches!(source, McpSourceSpec::File(McpFileSpec { defer_tools: true, .. }))); - } - #[test] fn serializes_deferred_mcp_file_as_typed_object() { let source: McpSourceSpec = McpFileSpec::new("mcp.json").defer_tools().into(); @@ -1104,7 +1057,7 @@ mod tests { #[test] fn load_default_resolves_workspace_scoped_user_prompt_and_mcp_paths() { - let project = project().file("AGENTS.md", "Agent instructions").file(".aether/mcp.json", r#"{"servers":{}}"#); + let project = project().file("AGENTS.md", "Agent instructions").file(".aether/mcp.json", &mcp_json("project")); let home = home().file(".aether/agents/planner/SYSTEM.md", "System instructions").settings( &json!({ "agents": [agent_json_with( @@ -1137,10 +1090,7 @@ mod tests { Prompt::File { path, .. } => path == &expected_agents, _ => false, })); - assert!(matches!( - &spec.mcp_config_sources[0], - McpConfigSource::File { path, defer_tools: false } if *path == project.root().join(".aether/mcp.json") - )); + assert_eq!(spec.mcp.servers.keys().collect::>(), ["project"]); } #[test] @@ -1214,7 +1164,7 @@ mod tests { let catalog = AgentCatalog::from_settings(project.root(), config).unwrap(); let spec = catalog.resolve("alpha").unwrap(); - assert!(spec.mcp_config_sources.is_empty()); + assert!(spec.mcp.is_empty()); } #[test] @@ -1233,7 +1183,7 @@ mod tests { let catalog = AgentCatalog::from_settings(project.root(), config).unwrap(); let spec = catalog.resolve("alpha").unwrap(); - assert!(spec.mcp_config_sources.is_empty()); + assert!(spec.mcp.is_empty()); } #[test] @@ -1270,7 +1220,7 @@ mod tests { #[test] fn optional_existing_mcp_source_preserves_defer_tools_flag() { - let project = project().file("BASE.md", "Base instructions").file("mcp.json", r#"{"servers":{}}"#); + let project = project().file("BASE.md", "Base instructions").file("mcp.json", &mcp_json("lazy")); let config = AetherSettings { agents: vec![AgentConfig { prompts: vec![PromptSource::file("BASE.md")], @@ -1283,7 +1233,7 @@ mod tests { let catalog = AgentCatalog::from_settings(project.root(), config).unwrap(); let spec = catalog.resolve("alpha").unwrap(); - assert!(matches!(&spec.mcp_config_sources[0], McpConfigSource::File { defer_tools: true, .. })); + assert!(spec.mcp.servers["lazy"].defer_tools().defers_tools()); } #[test] diff --git a/crates/aether-project/src/agent_catalog.rs b/crates/aether-project/src/agent_catalog.rs index dfd3cb21b..9b601be94 100644 --- a/crates/aether-project/src/agent_catalog.rs +++ b/crates/aether-project/src/agent_catalog.rs @@ -1,14 +1,13 @@ use crate::aether_settings::{project_settings_exist, user_settings_exist}; use crate::error::SettingsError; -use crate::{AetherSettings, AgentConfig, McpFileSpec, McpSourceSpec}; -use aether_core::agent_spec::{AgentSpec, AgentSpecExposure, McpConfigSource}; +use crate::{AetherSettings, AgentConfig, McpSourceSpec}; +use aether_core::agent_spec::{AgentSpec, AgentSpecExposure}; use aether_core::core::{AgentRegistry, Prompt}; +use aether_core::mcp::McpConfig; use llm::catalog::{LlmModel, ModelSpec, ModelSpecError}; use llm::{ProviderConnectionOverrides, ReasoningEffort}; -use mcp_utils::client::McpConfig; use std::collections::HashSet; use std::path::{Path, PathBuf}; -use utils::variables::VarError; /// A resolved catalog of agents from a project. #[derive(Debug, Clone)] @@ -217,7 +216,7 @@ fn resolve_agent_entry( }); } let mcp_sources = if entry.mcps.is_empty() { &defaults.mcps } else { &entry.mcps }; - let mcp_config_sources = resolve_mcp_config_sources(project_root, mcp_sources)?; + let mcp = load_mcp_config(project_root, mcp_sources)?; let mut provider_connections = defaults.providers.clone(); provider_connections.merge(entry.providers); @@ -230,44 +229,14 @@ fn resolve_agent_entry( context_window: entry.context_window, prompts, provider_connections, - mcp_config_sources, + mcp, exposure: AgentSpecExposure { user_invocable: entry.user_invocable, agent_invocable: entry.agent_invocable }, tools: entry.tools, }) } -fn resolve_mcp_config_sources( - workspace_root: &Path, - entries: &[McpSourceSpec], -) -> Result, SettingsError> { - entries - .iter() - .filter_map(|entry| match entry { - McpSourceSpec::File(McpFileSpec { path, defer_tools, optional }) => match path.resolve(workspace_root) { - Ok(full_path) => { - if full_path.is_file() { - Some(Ok(McpConfigSource::file(full_path, *defer_tools))) - } else if *optional { - None - } else { - Some(Err(SettingsError::InvalidMcpConfigPath { path: path.as_authored().to_string() })) - } - } - Err(VarError::NotFound(variable)) if *optional => { - tracing::warn!( - "Skipping optional MCP config '{}': variable '{variable}' is not defined", - path.as_authored() - ); - None - } - Err(VarError::NotFound(variable)) => Some(Err(SettingsError::UnresolvedMcpConfigVariable { - path: path.as_authored().to_string(), - variable, - })), - }, - McpSourceSpec::Inline { servers } => Some(Ok(McpConfigSource::Inline(McpConfig::new(servers.clone())))), - }) - .collect() +fn load_mcp_config(workspace_root: &Path, entries: &[McpSourceSpec]) -> Result { + entries.iter().filter_map(|entry| entry.load(workspace_root).transpose()).collect() } fn parse_model(agent: &str, model: &str) -> Result { @@ -281,7 +250,7 @@ fn parse_model(agent: &str, model: &str) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::testing::{DEFAULT_MODEL, agent_json_with, fake_spec, project, settings_agent}; + use crate::testing::{DEFAULT_MODEL, agent_json_with, fake_spec, mcp_json, project, settings_agent}; use aether_core::agent_spec::AgentSpecExposure; use llm::ModelSettings; use serde_json::json; @@ -291,14 +260,8 @@ mod tests { AgentCatalog::new(project_root, vec![planner], None) } - fn file_sources(spec: &AgentSpec) -> Vec<(PathBuf, bool)> { - spec.mcp_config_sources - .iter() - .filter_map(|source| match source { - McpConfigSource::File { path, defer_tools } => Some((path.clone(), *defer_tools)), - McpConfigSource::Json(_) | McpConfigSource::Inline(_) => None, - }) - .collect() + fn mcp_servers(spec: &AgentSpec) -> Vec<&str> { + spec.mcp.servers.keys().map(String::as_str).collect() } fn has_prompt_file(spec: &AgentSpec, expected: &str) -> bool { @@ -596,7 +559,7 @@ mod tests { #[test] fn top_level_mcps_are_inherited_when_agent_mcps_are_empty() { - let dir = project().file("BASE.md", "Base instructions").file("base-mcp.json", "{}"); + let dir = project().file("BASE.md", "Base instructions").file("base-mcp.json", &mcp_json("base")); let config = AetherSettings { prompts: vec![crate::PromptSource::file("BASE.md")], @@ -608,13 +571,15 @@ mod tests { let catalog = AgentCatalog::from_settings(dir.root(), config).unwrap(); let spec = catalog.resolve("planner").unwrap(); - assert_eq!(file_sources(&spec), vec![(dir.root().join("base-mcp.json"), false)]); + assert_eq!(mcp_servers(&spec), ["base"]); } #[test] fn agent_mcps_override_top_level_mcps() { - let dir = - project().file("BASE.md", "Base instructions").file("base-mcp.json", "{}").file("agent-mcp.json", "{}"); + let dir = project() + .file("BASE.md", "Base instructions") + .file("base-mcp.json", &mcp_json("base")) + .file("agent-mcp.json", &mcp_json("agent")); let config = AetherSettings { prompts: vec![crate::PromptSource::file("BASE.md")], @@ -629,7 +594,22 @@ mod tests { let catalog = AgentCatalog::from_settings(dir.root(), config).unwrap(); let spec = catalog.resolve("planner").unwrap(); - assert_eq!(file_sources(&spec), vec![(dir.root().join("agent-mcp.json"), false)]); + assert_eq!(mcp_servers(&spec), ["agent"]); + } + + #[test] + fn unreadable_mcp_config_names_its_path() { + let dir = project().file("BASE.md", "Base instructions").file("broken.json", "not json"); + let config = AetherSettings { + prompts: vec![crate::PromptSource::file("BASE.md")], + mcps: vec![McpSourceSpec::file("broken.json")], + agents: vec![settings_agent("planner", "Planner agent")], + ..AetherSettings::default() + }; + + let error = AgentCatalog::from_settings(dir.root(), config).unwrap_err(); + + assert!(matches!(error, SettingsError::McpConfig { ref path, .. } if path == "broken.json"), "{error}"); } #[test] @@ -652,15 +632,14 @@ mod tests { #[test] fn resolve_preserves_agent_mcp() { - let dir = project().file("agent-mcp.json", "{}"); - + let dir = project(); let mut planner = fake_spec("planner", AgentSpecExposure::both()); - planner.mcp_config_sources = vec![McpConfigSource::model_visible(dir.root().join("agent-mcp.json"))]; + planner.mcp = McpConfig::from_json(&mcp_json("agent")).unwrap(); let catalog = AgentCatalog::new(dir.root().to_path_buf(), vec![planner], None); let spec = catalog.resolve("planner").unwrap(); - assert_eq!(file_sources(&spec), vec![(dir.root().join("agent-mcp.json"), false)]); + assert_eq!(mcp_servers(&spec), ["agent"]); } #[test] @@ -670,6 +649,6 @@ mod tests { AgentCatalog::new(dir.root().to_path_buf(), vec![fake_spec("planner", AgentSpecExposure::both())], None); let spec = catalog.resolve("planner").unwrap(); - assert!(spec.mcp_config_sources.is_empty()); + assert!(spec.mcp.is_empty()); } } diff --git a/crates/aether-project/src/agent_config.rs b/crates/aether-project/src/agent_config.rs index 1f395e868..4f77b90d6 100644 --- a/crates/aether-project/src/agent_config.rs +++ b/crates/aether-project/src/agent_config.rs @@ -1,6 +1,6 @@ use crate::{McpSourceSpec, PromptSource}; +use aether_core::mcp::ToolFilter; use llm::{ModelSettings, ProviderConnectionOverrides, ReasoningEffort}; -use mcp_utils::client::ToolFilter; #[doc = include_str!("docs/agent_config.md")] #[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] diff --git a/crates/aether-project/src/error.rs b/crates/aether-project/src/error.rs index 2da862228..9751e6087 100644 --- a/crates/aether-project/src/error.rs +++ b/crates/aether-project/src/error.rs @@ -74,6 +74,13 @@ pub enum SettingsError { #[error("MCP config path '{path}' references undefined variable '{variable}'")] UnresolvedMcpConfigVariable { path: String, variable: String }, + #[error("Failed to read MCP config '{path}': {source}")] + McpConfig { + path: String, + #[source] + source: aether_core::mcp::ParseError, + }, + /// I/O error while reading files. #[error("I/O error: {0}")] IoError(String), diff --git a/crates/aether-project/src/lib.rs b/crates/aether-project/src/lib.rs index 331104ae3..b53bb5299 100644 --- a/crates/aether-project/src/lib.rs +++ b/crates/aether-project/src/lib.rs @@ -4,7 +4,7 @@ pub mod aether_settings; mod agent_catalog; mod agent_config; mod error; -mod mcp_config_source_config; +mod mcp_source_spec; mod prompt_catalog; pub mod prompt_file; #[cfg(feature = "testing")] @@ -19,6 +19,6 @@ pub use aether_settings::{ pub use agent_catalog::AgentCatalog; pub use agent_config::AgentConfig; pub use error::SettingsError; -pub use mcp_config_source_config::{McpFileSpec, McpSourceSpec}; +pub use mcp_source_spec::{McpFileSpec, McpSourceSpec}; pub use prompt_catalog::PromptCatalog; pub use prompt_file::{PromptFile, PromptFileError, PromptTriggers, SKILL_FILENAME}; diff --git a/crates/aether-project/src/mcp_config_source_config.rs b/crates/aether-project/src/mcp_source_spec.rs similarity index 59% rename from crates/aether-project/src/mcp_config_source_config.rs rename to crates/aether-project/src/mcp_source_spec.rs index fa46257ca..8afdfb48c 100644 --- a/crates/aether-project/src/mcp_config_source_config.rs +++ b/crates/aether-project/src/mcp_source_spec.rs @@ -1,12 +1,15 @@ -use mcp_utils::client::McpServerConfig; +use crate::SettingsError; +use aether_core::mcp::{McpConfig, McpServerConfig}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; +use utils::variables::VarError; use utils::{PathOrObject, ResourcePath, is_false, string_or_object_schema}; #[derive(Debug, Clone, PartialEq)] pub enum McpSourceSpec { File(McpFileSpec), - Inline { servers: BTreeMap }, + Inline(McpConfig), } #[derive(Debug, Clone, PartialEq)] @@ -32,6 +35,43 @@ impl McpFileSpec { self.optional = true; self } + + pub fn load(&self, root: &Path) -> Result, SettingsError> { + let Some(path) = self.resolve(root)? else { return Ok(None) }; + let mut config = McpConfig::from_json_file(&path) + .map_err(|source| SettingsError::McpConfig { path: self.path.as_authored().to_string(), source })?; + if self.defer_tools { + config.defer_all_tools(); + } + Ok(Some(config)) + } + + fn resolve(&self, root: &Path) -> Result, SettingsError> { + let full_path = match self.path.resolve(root) { + Ok(full_path) => full_path, + Err(VarError::NotFound(variable)) if self.optional => { + tracing::warn!( + "Skipping optional MCP config '{}': variable '{variable}' is not defined", + self.path.as_authored() + ); + return Ok(None); + } + Err(VarError::NotFound(variable)) => { + return Err(SettingsError::UnresolvedMcpConfigVariable { + path: self.path.as_authored().to_string(), + variable, + }); + } + }; + + if full_path.is_file() { + Ok(Some(full_path)) + } else if self.optional { + Ok(None) + } else { + Err(SettingsError::InvalidMcpConfigPath { path: self.path.as_authored().to_string() }) + } + } } impl McpSourceSpec { @@ -42,7 +82,14 @@ impl McpSourceSpec { pub fn path(&self) -> Option<&str> { match self { Self::File(file) => Some(file.path.as_authored()), - Self::Inline { .. } => None, + Self::Inline(_) => None, + } + } + + pub fn load(&self, root: &Path) -> Result, SettingsError> { + match self { + Self::File(file) => file.load(root), + Self::Inline(config) => Ok(Some(config.clone())), } } } @@ -60,7 +107,7 @@ impl<'de> Deserialize<'de> for McpSourceSpec { PathOrObject::Object(McpSourceSpecObject::File { path, defer_tools, optional }) => { Self::File(McpFileSpec { path, defer_tools, optional }) } - PathOrObject::Object(McpSourceSpecObject::Inline { servers }) => Self::Inline { servers }, + PathOrObject::Object(McpSourceSpecObject::Inline { servers }) => Self::Inline(McpConfig::new(servers)), }) } } @@ -73,8 +120,8 @@ impl Serialize for McpSourceSpec { &McpSourceSpecObject::File { path: path.clone(), defer_tools: *defer_tools, optional: *optional }, serializer, ), - Self::Inline { servers } => { - Serialize::serialize(&McpSourceSpecObject::Inline { servers: servers.clone() }, serializer) + Self::Inline(config) => { + Serialize::serialize(&McpSourceSpecObject::Inline { servers: config.servers.clone() }, serializer) } } } @@ -98,7 +145,7 @@ impl schemars::JsonSchema for McpSourceSpec { enum McpSourceSpecObject { File { path: ResourcePath, - #[serde(rename = "deferTools", alias = "proxy", default, skip_serializing_if = "is_false")] + #[serde(rename = "deferTools", default, skip_serializing_if = "is_false")] defer_tools: bool, #[serde(default, skip_serializing_if = "is_false")] optional: bool, diff --git a/crates/aether-project/src/testing.rs b/crates/aether-project/src/testing.rs index 3a801fd89..bde7dcb50 100644 --- a/crates/aether-project/src/testing.rs +++ b/crates/aether-project/src/testing.rs @@ -12,13 +12,17 @@ use tempfile::TempDir; use crate::prompt_file::{PromptFile, PromptTriggers}; use crate::{AgentConfig, PromptSource, SKILL_FILENAME}; use aether_core::agent_spec::{AgentSpec, AgentSpecExposure}; +use aether_core::mcp::{McpConfig, ToolFilter}; use llm::{ModelSettings, ProviderConnectionOverrides}; -use mcp_utils::client::ToolFilter; /// The model id every shared agent fixture resolves to. pub const DEFAULT_MODEL: &str = "anthropic:claude-sonnet-4-5"; /// A temporary project directory for tests that need on-disk settings, prompts, or skills. +pub fn mcp_json(server: &str) -> String { + json!({ "servers": { server: { "command": format!("{server}-server") } } }).to_string() +} + pub fn project() -> TestProject { TestProject { root: temp_dir("project") } } @@ -98,7 +102,7 @@ pub fn fake_spec(name: &str, exposure: AgentSpecExposure) -> AgentSpec { context_window: None, prompts: vec![], provider_connections: ProviderConnectionOverrides::default(), - mcp_config_sources: Vec::new(), + mcp: McpConfig::default(), exposure, tools: ToolFilter::default(), } diff --git a/crates/aether-sessions/tests/message_identity.rs b/crates/aether-sessions/tests/message_identity.rs index 2d8385dcf..15514c044 100644 --- a/crates/aether-sessions/tests/message_identity.rs +++ b/crates/aether-sessions/tests/message_identity.rs @@ -26,13 +26,12 @@ async fn live_context_and_reloaded_log_share_identities_and_message_boundaries() let (_, loaded) = store.store().load("identity").unwrap(); let restored = context_from_events(&loaded); let replayed = context_from_events(&loaded); - let ids = |messages: &[ChatMessage]| messages.iter().map(ChatMessage::message_id).collect::>(); - assert_eq!(ids(restored.messages()), ids(replayed.messages())); + assert_eq!(conversation_ids(restored.messages()), conversation_ids(replayed.messages())); assert_eq!(restored.message_count(), 4); assert_eq!(restored.messages()[0].message_id(), Some(user_id)); let contexts = result.captured_contexts.lock().unwrap(); - assert_eq!(ids(contexts[1].messages()), ids(&restored.messages()[..3])); + assert_eq!(conversation_ids(contexts[1].messages()), conversation_ids(&restored.messages()[..3])); let assistant_ids: Vec<_> = restored .messages() .iter() @@ -59,7 +58,7 @@ async fn live_context_and_reloaded_log_share_identities_and_message_boundaries() } let mut switched = contexts[1].clone(); switched.replace_conversation(restored.messages().clone()); - assert_eq!(ids(switched.messages()), ids(restored.messages())); + assert_eq!(conversation_ids(switched.messages()), conversation_ids(restored.messages())); } #[tokio::test] @@ -77,15 +76,12 @@ async fn continuation_preserves_live_message_boundaries_and_ids() { .await .unwrap(); let contexts = result.captured_contexts.lock().unwrap(); - assert_eq!(contexts[1].message_count(), 3); let mut events = vec![SessionEvent::User(UserEvent::Message { message_id: user_id, content, display_content: None })]; events.extend(result.messages.into_iter().map(SessionEvent::Agent)); let restored = context_from_events(&events); assert_eq!(restored.message_count(), 4); - for (live, restored) in contexts[1].messages().iter().zip(restored.messages()) { - assert_eq!(live.message_id(), restored.message_id()); - } + assert_eq!(conversation_ids(contexts[1].messages()), conversation_ids(&restored.messages()[..3])); } #[test] @@ -127,3 +123,7 @@ fn partial_streams_do_not_replace_durable_messages() { assert_eq!(context.message_count(), 1); assert_eq!(context.messages()[0].message_id().as_deref(), Some("original")); } + +fn conversation_ids(messages: &[ChatMessage]) -> Vec> { + messages.iter().filter(|message| !message.is_system()).map(ChatMessage::message_id).collect() +} diff --git a/crates/mcp-servers/Cargo.toml b/crates/mcp-servers/Cargo.toml index 8525592f9..52fce6f1a 100644 --- a/crates/mcp-servers/Cargo.toml +++ b/crates/mcp-servers/Cargo.toml @@ -27,14 +27,13 @@ workspace = true [features] default = ["coding", "skills", "tasks", "subagents", "review"] test-helpers = ["dep:tempfile"] -file-ops = ["dep:thiserror"] +file-ops = [] coding = [ "file-ops", "dep:aether-project", "dep:aether-lspd", "dep:futures", "dep:tracing", - "dep:thiserror", "dep:lsp-types", "dep:reqwest", "dep:chrono", @@ -42,16 +41,12 @@ coding = [ "dep:dom_smoothie", "dep:url", ] -skills = [ - "dep:aether-project", - "dep:tracing", -] +skills = ["dep:aether-project"] tasks = [ "dep:chrono", "dep:tempfile", - "dep:thiserror", ] -review = ["dep:thiserror", "dep:axum", "dep:uuid", "dep:reqwest", "dep:tokio-util", "dep:tower-http", "tokio/net"] +review = ["dep:axum", "dep:uuid", "dep:reqwest", "dep:tokio-util", "dep:tower-http", "tokio/net"] subagents = [ "coding", "skills", @@ -62,10 +57,7 @@ subagents = [ "dep:aether-project", "dep:llm", "dep:futures", - "dep:tracing", ] -all = ["coding", "skills", "tasks", "subagents", "review"] -stdio = ["coding", "skills", "tasks", "subagents", "review", "rmcp/transport-io"] [dependencies] # Always required @@ -76,12 +68,12 @@ libc = { workspace = true } rmcp = { workspace = true, features = ["elicitation", "macros", "server"] } schemars = { workspace = true } clap = { workspace = true, features = ["env"] } +thiserror = { workspace = true } mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", default-features = false, version = "0.7.13" } utils = { package = "aether-utils", path = "../utils", version = "0.2.19" } aether-lspd = { path = "../aether-lspd", optional = true, version = "0.3.3" } futures = { workspace = true, optional = true } tracing = { workspace = true, optional = true } -thiserror = { workspace = true, optional = true } lsp-types = { workspace = true, optional = true } reqwest = { workspace = true, optional = true } chrono = { workspace = true, optional = true } @@ -109,12 +101,13 @@ reqwest = { workspace = true } tokio = { workspace = true, features = ["test-util"] } aether-lspd = { path = "../aether-lspd", features = ["testing"], version = "0.3.3" } mcp-servers = { package = "aether-mcp-servers", path = ".", features = ["test-helpers"] } -mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", features = ["client"], version = "0.7.13" } +mcp_utils = { package = "aether-mcp-utils", path = "../mcp-utils", features = ["testing"], version = "0.7.13" } +aether-core = { package = "aether-agent-core", path = "../aether-core", features = ["testing"], version = "0.9.4" } [[bin]] name = "mcp-servers-stdio" path = "src/bin/stdio.rs" -required-features = ["stdio"] +required-features = ["subagents"] [[example]] name = "lsp_diagnostics" diff --git a/crates/mcp-servers/README.md b/crates/mcp-servers/README.md index 3ffb515b4..e93c596fb 100644 --- a/crates/mcp-servers/README.md +++ b/crates/mcp-servers/README.md @@ -66,16 +66,17 @@ These servers use Aether's `in-memory` transport type -- they run inside your ag } ``` -Each server key must match a factory registered with `McpBuilder::register_in_memory_server()`. Config loading only records a cloneable in-memory server specification and expands its variables. The factory is invoked exactly once during `McpBuilder::spawn`, after the runtime MCP handle exists. It receives the specification plus `RuntimeServices` containing the live `McpHandle`, root directory, and `AgentDeps`. The `args` array is parsed as CLI flags by each server: +Each server key must match a factory registered with `McpBuilder::register_in_memory_server()`. Loading the config only expands the server's variables. The factory is invoked exactly once during `McpBuilder::spawn`, after the deferred-tools endpoint is bound. It receives the server's `args`, with variables expanded, plus `RuntimeServices` containing the root directory, `AgentDeps`, and the deferred-tools socket path (when any server defers its tools). The `args` array is parsed as CLI flags by each server, with relative paths resolved against the root directory. Invalid `args` fail `spawn` with `McpSpawnError::ServerFactory` rather than falling back to defaults: | Server | Flag | Default | Description | |--------|------|---------|-------------| -| `coding` | `--root-dir ` | cwd | Workspace root for LSP and file operations | +| `coding` | `--root-dir ` | workspace root | Workspace root for LSP and file operations | | `coding` | `--rules-dir ` (repeatable) | none | Explicit prompt directories for automatic read-triggered rules | +| `coding` | `--permission-mode ` | `always-allow` | `always-allow`, `auto`, or `always-ask` | | `coding` | `--disable-lsp` | enabled | Disable LSP-backed tools and `aether-lspd` daemon connections | | `skills` | `--dir ` (repeatable) | required | Prompt directories to scan | -| `tasks` | `--dir ` | `.` | Base directory for task storage (creates `.aether-tasks/` inside) | -| `subagents` | `--project-root ` (alias: `--dir`) | `.` | Project root containing optional `.aether/settings.json` authored agents | +| `tasks` | `--dir ` | session-scoped temp dir | Base directory for persistent task storage (creates `.aether-tasks/` inside) | +| `subagents` | `--project-root ` (alias: `--dir`) | workspace root | Project root containing optional `.aether/settings.json` authored agents | | `review` | `--root-dir ` | workspace root | Base for relative artifact paths | To register factories and load the config: @@ -83,25 +84,23 @@ To register factories and load the config: ```rust,ignore use aether_core::mcp::mcp; use mcp_servers::McpBuilderExt; +use aether_core::mcp::McpConfig; let builder = mcp("/my/project") .with_agent_deps(deps) .with_builtin_servers() - .from_json_files(&["mcp.json"])?; + .with_config(McpConfig::from_json_file("mcp.json")?)?; // Built-in factories have not run yet. They receive RuntimeServices here. -let runtime = builder.spawn().await?; +let runtime = builder.spawn()?; ``` -Custom factories use the same lazy signature: +Custom factories use the same lazy signature and return an error for invalid `args`: ```rust,ignore -builder.register_in_memory_server("custom", Box::new(|spec, services| { - async move { - CustomMcp::new(spec.args, services.root_dir, services.mcp).into_dyn() - } - .boxed() -})); +builder.register_in_memory_server("custom", |args, services| { + Ok(McpServer::new(CustomMcp::from_args(args, &services.root_dir)?)) +}); ``` ## Programmatic Usage @@ -120,18 +119,15 @@ Create and start servers directly: ```rust,ignore use mcp_servers::{CodingMcp, SkillsMcp, TasksMcp}; -use rmcp::ServiceExt; +use mcp_utils::server::McpServer; // Create a coding server -let server = CodingMcp::new() - .with_root_dir("/my/project".into()) - .into_dyn(); +let server = McpServer::new(CodingMcp::new("/my/project")); // Or with LSP support -let server = CodingMcp::new() - .with_root_dir("/my/project".into()) - .with_lsp("/my/project".into()) - .into_dyn(); +let server = McpServer::new(CodingMcp::new("/my/project").with_lsp()); + +server.serve_stdio().await?; ``` ## Server Documentation diff --git a/crates/mcp-servers/src/background_tasks.rs b/crates/mcp-servers/src/background_tasks.rs new file mode 100644 index 000000000..496bc9a5b --- /dev/null +++ b/crates/mcp-servers/src/background_tasks.rs @@ -0,0 +1,80 @@ +use rmcp::ErrorData; +use rmcp::model::{CallToolResult, CancelTaskParams, GetTaskParams, GetTaskResult, Task, UpdateTaskParams}; +use rmcp::task_manager::{TaskContext, TaskExit, TaskManager, TaskOptions}; +use std::future::Future; + +pub(crate) struct BackgroundTasks { + manager: TaskManager, +} + +impl BackgroundTasks { + pub(crate) fn new() -> Self { + Self { manager: TaskManager::new() } + } + + pub(crate) fn spawn( + &self, + status: impl Into, + work: impl Future> + Send + 'static, + ) -> Task { + let options = TaskOptions::new().with_ttl_ms(BACKGROUND_TASK_TTL_MS).with_status_message(status); + self.manager.spawn(options, move |context: TaskContext| { + Box::pin(async move { + tokio::select! { + () = context.cancelled() => Err(TaskExit::Cancelled), + result = work => result, + } + }) + }) + } + + pub(crate) fn get(&self, request: &GetTaskParams) -> Result { + self.manager.get_task(&request.task_id).map(GetTaskResult::new) + } + + pub(crate) fn update(&self, request: UpdateTaskParams) -> Result<(), ErrorData> { + self.manager.update_task(&request.task_id, request.input_responses) + } + + pub(crate) fn cancel(&self, request: &CancelTaskParams) -> Result<(), ErrorData> { + self.manager.cancel_task(&request.task_id) + } +} + +impl Drop for BackgroundTasks { + fn drop(&mut self) { + self.manager.shutdown(); + } +} + +macro_rules! delegate_background_tasks { + ($tasks:ident) => { + fn get_task( + &self, + request: rmcp::model::GetTaskParams, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + Send + '_ { + std::future::ready(self.$tasks.get(&request)) + } + + fn update_task( + &self, + request: rmcp::model::UpdateTaskParams, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + Send + '_ { + std::future::ready(self.$tasks.update(request)) + } + + fn cancel_task( + &self, + request: rmcp::model::CancelTaskParams, + _context: rmcp::service::RequestContext, + ) -> impl std::future::Future> + Send + '_ { + std::future::ready(self.$tasks.cancel(&request)) + } + }; +} + +pub(crate) use delegate_background_tasks; + +const BACKGROUND_TASK_TTL_MS: u64 = 3_600_000; diff --git a/crates/mcp-servers/src/bin/stdio.rs b/crates/mcp-servers/src/bin/stdio.rs index f31bcd943..f50ac5b19 100644 --- a/crates/mcp-servers/src/bin/stdio.rs +++ b/crates/mcp-servers/src/bin/stdio.rs @@ -1,16 +1,15 @@ +use aether_core::core::AgentDeps; +use aether_core::mcp::RuntimeServices; use clap::Parser; -use mcp_servers::workspace_paths::{current_dir, resolve_path}; -use mcp_servers::{CodingMcp, CodingMcpArgs, ReviewMcp, SkillsMcp, SubAgentsMcp, TasksMcp}; -use mcp_utils::ServiceExt; -use rmcp::ServerHandler; -use rmcp::transport::io::stdio; +use mcp_servers::workspace_paths::current_dir; +use mcp_servers::{Builtin, SubAgentsMcp, SubAgentsMcpArgs}; +use mcp_utils::server::{McpServer, ServeError}; #[derive(Parser)] #[command(name = "mcp-servers-stdio", about = "Run an MCP server over stdio")] struct Cli { - /// Which server to run: coding, skills, tasks, subagents, review - #[arg(long)] - server: String, + #[arg(long, value_enum)] + server: Builtin, /// Arguments forwarded to the selected server (e.g. --root-dir /path) #[arg(trailing_var_arg = true, allow_hyphen_values = true)] @@ -19,55 +18,23 @@ struct Cli { #[derive(Debug, thiserror::Error)] enum StdioError { - #[error("Unknown server: '{0}'. Available: coding, skills, tasks, subagents, review")] - UnknownServer(String), - #[error("{0}")] - ServerArgs(#[from] mcp_servers::error::ServerInitError), - #[error("Failed to start server: {0}")] - Serve(String), - #[error("Server task failed: {0}")] - Join(tokio::task::JoinError), -} - -async fn serve_stdio(server: impl ServerHandler) -> Result<(), StdioError> { - let running = server.serve(stdio()).await.map_err(|e| StdioError::Serve(e.to_string()))?; - running.waiting().await.map_err(StdioError::Join)?; - Ok(()) + #[error(transparent)] + Settings(#[from] aether_project::SettingsError), + #[error(transparent)] + Serve(#[from] ServeError), } #[tokio::main] async fn main() -> Result<(), StdioError> { - let cli = Cli::parse(); - - match cli.server.as_str() { - "coding" => { - let CodingMcpArgs { root_dir, rules_dirs, permission_mode, disable_lsp } = - CodingMcpArgs::from_args(cli.args).map_err(StdioError::ServerArgs)?; - let root_dir = root_dir.unwrap_or_else(current_dir); - let rules_dirs = rules_dirs.into_iter().map(|path| resolve_path(&root_dir, path)).collect(); - let server = CodingMcp::new() - .with_root_dir(root_dir.clone()) - .with_rules_dirs(rules_dirs) - .with_permission_mode(permission_mode); - let server = if disable_lsp { server } else { server.with_lsp(root_dir) }; - serve_stdio(server).await - } - "skills" => { - let server = SkillsMcp::from_args(cli.args).map_err(StdioError::ServerArgs)?; - serve_stdio(server).await - } - "tasks" => { - let server = TasksMcp::from_args(cli.args).map_err(StdioError::ServerArgs)?; - serve_stdio(server).await - } - "subagents" => { - let server = SubAgentsMcp::standalone_from_args(cli.args).map_err(StdioError::ServerArgs)?; - serve_stdio(server).await - } - "review" => { - let server = ReviewMcp::from_args(cli.args).map_err(StdioError::ServerArgs)?; - serve_stdio(server).await + let Cli { server, args } = Cli::parse(); + let services = + RuntimeServices { root_dir: current_dir(), agent_deps: AgentDeps::default(), deferred_tools_socket: None }; + let server: McpServer = match server { + Builtin::Subagents => { + let args = SubAgentsMcpArgs::try_parse_from(args).unwrap_or_else(|error| error.exit()); + SubAgentsMcp::standalone(args, &services.root_dir)?.into() } - other => Err(StdioError::UnknownServer(other.to_string())), - } + builtin => builtin.build(args, &services).unwrap_or_else(|error| error.exit()), + }; + Ok(server.serve_stdio().await?) } diff --git a/crates/mcp-servers/src/coding/mod.rs b/crates/mcp-servers/src/coding/mod.rs index 52d7a84b2..126d91522 100644 --- a/crates/mcp-servers/src/coding/mod.rs +++ b/crates/mcp-servers/src/coding/mod.rs @@ -1,3 +1,4 @@ +use crate::server_config; use aether_project::PromptCatalog; use clap::Parser; use rmcp::{ @@ -8,18 +9,16 @@ use rmcp::{ wrapper::{Json, Parameters}, }, model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ContentBlock, CreateTaskResult, - ElicitRequest, ElicitRequestParams, ElicitResult, ElicitationAction, ElicitationSchema, EnumSchema, - GetTaskParams, GetTaskResult, Implementation, InputRequest, InputRequests, InputRequiredResult, - ProgressNotificationParam, ServerCapabilities, ServerConfig, UpdateTaskParams, + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, CreateTaskResult, ElicitRequest, + ElicitRequestParams, ElicitResult, ElicitationAction, ElicitationSchema, EnumSchema, InputRequest, + InputRequests, InputRequiredResult, ProgressNotificationParam, ServerCapabilities, ServerConfig, }, - service::RequestContext, - task_manager::{TaskContext, TaskExit, TaskManager, TaskOptions}, + service::{ElicitationMode, RequestContext}, tool, tool_handler, tool_router, }; use std::collections::HashSet; use std::fmt::{Debug, Formatter, Write as _}; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use tokio::fs::try_exists; @@ -32,17 +31,15 @@ pub mod tools_trait; pub use default_tools::DefaultCodingTools; pub use tools_trait::CodingTools; +use crate::background_tasks::{BackgroundTasks, delegate_background_tasks}; use crate::coding::error::CodingError; +use crate::coding::prompt_rule_matcher::PromptRuleMatcher; +use crate::lsp::registry::LspRegistry; use crate::lsp::tools::check_errors::{LspDiagnosticsOutput, LspDiagnosticsRequest, execute_lsp_diagnostics}; +use crate::lsp::tools::rename::{LspRenameInput, LspRenameOutput, execute_lsp_rename}; use crate::lsp::tools::symbol_lookup::{LspSymbolInput, LspSymbolOutput, execute_lsp_symbol}; -use crate::workspace_paths::WorkspacePaths; -use crate::{coding::prompt_rule_matcher::PromptRuleMatcher, lsp::registry::LspRegistry}; -use crate::{ - error::ServerInitError, - lsp::tools::rename::{LspRenameInput, LspRenameOutput, execute_lsp_rename}, -}; -use mcp_utils::server::mrtr::{input_requests_supported, parse_response}; -use mcp_utils::server::tasks::{BACKGROUND_TASK_TTL_MS, require_tasks_capability}; +use crate::request_context::{CLIENT_CANNOT_ELICIT, parse_response, require_tasks, supports_elicitation}; +use crate::workspace_paths::{WorkspacePaths, resolve_path, resolve_root}; use tools::bash::{BashInput, BashOutput, validate_args}; use tools::edit_file::{EditFileArgs, EditFileResponse, edit_file_contents}; @@ -68,7 +65,8 @@ pub enum PermissionMode { } /// CLI arguments for `CodingMcp` server -#[derive(Debug, Clone, Default, Parser)] +#[derive(Debug, Clone, Parser)] +#[command(name = "coding", no_binary_name = true)] pub struct CodingMcpArgs { /// Root directory for path resolution and LSP initialization. #[arg(long = "root-dir")] @@ -88,19 +86,10 @@ pub struct CodingMcpArgs { pub disable_lsp: bool, } -impl CodingMcpArgs { - pub fn from_args(args: Vec) -> Result { - let mut full_args = vec!["coding-mcp".to_string()]; - full_args.extend(args); - - Self::try_parse_from(full_args).map_err(ServerInitError::InvalidArgs) - } -} - #[doc = include_str!("../docs/coding_mcp.md")] pub struct CodingMcp { tool_router: ToolRouter, - task_manager: TaskManager, + tasks: BackgroundTasks, read_state: Mutex, tools: Arc, /// Optional LSP operations (enabled with `.with_lsp()`) @@ -122,10 +111,11 @@ pub struct CodingMcp { #[tool_handler(router = self.tool_router)] impl ServerHandler for CodingMcp { fn get_info(&self) -> ServerConfig { - let instructions = self.build_instructions(); - ServerConfig::new(ServerCapabilities::builder().enable_tools().enable_tasks().build()) - .with_server_info(Implementation::new("coding-mcp", "0.1.0")) - .with_instructions(instructions) + server_config( + "coding-mcp", + ServerCapabilities::builder().enable_tools().enable_tasks().build(), + self.build_instructions(), + ) } async fn call_tool( @@ -135,12 +125,16 @@ impl ServerHandler for CodingMcp { ) -> Result { let background_args = background_bash_args(&request); if let Some(args) = &background_args { - require_tasks_capability(&context)?; + require_tasks(&context)?; validate_args(args).map_err(|error| ErrorData::invalid_params(error.to_string(), None))?; } if let Some(prompt) = self.permission_prompt(&request) { let Some(responses) = request.input_responses.clone() else { + if !supports_elicitation(&context, ElicitationMode::Form) { + let message = permission_unsupported(&prompt.tool_name); + return Ok(CallToolResult::error(vec![ContentBlock::text(message)]).into()); + } let requests = InputRequests::from([( "permission".to_string(), InputRequest::Elicitation(ElicitRequest::new(decision_form( @@ -148,10 +142,6 @@ impl ServerHandler for CodingMcp { &prompt.description, ))), )]); - if !input_requests_supported(context.client_capabilities().as_ref(), &requests) { - let message = permission_unsupported(&prompt.tool_name); - return Ok(CallToolResult::error(vec![ContentBlock::text(message)]).into()); - } return Ok(InputRequiredResult::from_input_requests(requests).into()); }; @@ -171,35 +161,13 @@ impl ServerHandler for CodingMcp { self.tool_router.call(ToolCallContext::new(self, request, context)).await } - fn get_task( - &self, - request: GetTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.get_task(&request.task_id).map(GetTaskResult::new)) - } - - fn update_task( - &self, - request: UpdateTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.update_task(&request.task_id, request.input_responses)) - } - - fn cancel_task( - &self, - request: CancelTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.cancel_task(&request.task_id)) - } + delegate_background_tasks!(tasks); } impl CodingMcp { /// Create a new `CodingMcp` with default (local filesystem) tools - pub fn new() -> Self { - Self::with_tools(DefaultCodingTools::new()) + pub fn new(root_dir: impl Into) -> Self { + Self::with_tools(root_dir, DefaultCodingTools::new()) } } @@ -252,16 +220,16 @@ fn is_dangerous_cmd(command: &str) -> bool { #[tool_router] impl CodingMcp { /// Create a `CodingMcp` with custom tool implementation - pub fn with_tools(tools: T) -> Self { + pub fn with_tools(root_dir: impl Into, tools: T) -> Self { Self { tool_router: Self::tool_router(), - task_manager: TaskManager::new(), + tasks: BackgroundTasks::new(), read_state: Mutex::new(ReadState::default()), tools: Arc::new(tools), lsp: None, web_fetcher: WebFetcher::new(), web_searcher: WebSearcher::try_new().ok(), - root_dir: crate::workspace_paths::current_dir(), + root_dir: root_dir.into(), configured_rules_dirs: Vec::new(), permission_mode: PermissionMode::AlwaysAllow, spill_dir: Arc::new(TempDir::new()), @@ -270,18 +238,21 @@ impl CodingMcp { } } - /// Enable LSP code intelligence for the given project root. + pub fn from_args(args: CodingMcpArgs, base_dir: &Path, tools: T) -> Self { + let CodingMcpArgs { root_dir, rules_dirs, permission_mode, disable_lsp } = args; + let root_dir = resolve_root(base_dir, root_dir); + let rules_dirs = rules_dirs.into_iter().map(|path| resolve_path(&root_dir, path)).collect(); + let server = + Self::with_tools(root_dir, tools).with_rules_dirs(rules_dirs).with_permission_mode(permission_mode); + if disable_lsp { server } else { server.with_lsp() } + } + + /// Enable LSP code intelligence for the server's root directory. /// /// LSP servers for detected project languages are spawned immediately /// in the background, allowing indexing to start right away. - pub fn with_lsp(mut self, root_path: PathBuf) -> Self { - self.lsp = Some(LspRegistry::new_and_spawn(root_path)); - self - } - - /// Set the root directory. - pub fn with_root_dir(mut self, root_dir: PathBuf) -> Self { - self.root_dir = root_dir; + pub fn with_lsp(mut self) -> Self { + self.lsp = Some(LspRegistry::new_and_spawn(self.root_dir.clone())); self } @@ -305,20 +276,14 @@ impl CodingMcp { let spill_dir = Arc::clone(&self.spill_dir); let truncator = self.bash_truncator; let description = args.description.clone().unwrap_or_else(|| truncate(&args.command, 80)); - let options = TaskOptions::new().with_ttl_ms(BACKGROUND_TASK_TTL_MS).with_status_message(description); - - self.task_manager.spawn(options, move |task_context: TaskContext| { - Box::pin(async move { - tokio::select! { - () = task_context.cancelled() => Err(TaskExit::Cancelled), - result = run_bash(tools.as_ref(), args, cwd, truncator, &spill_dir) => Ok(match result { - Ok(output) => serde_json::to_value(output).map_or_else( - |error| CallToolResult::error(vec![ContentBlock::text(error.to_string())]), - CallToolResult::structured, - ), - Err(error) => CallToolResult::error(vec![ContentBlock::text(error.to_string())]), - }), - } + + self.tasks.spawn(description, async move { + Ok(match run_bash(tools.as_ref(), args, cwd, truncator, &spill_dir).await { + Ok(output) => serde_json::to_value(output).map_or_else( + |error| CallToolResult::error(vec![ContentBlock::text(error.to_string())]), + CallToolResult::structured, + ), + Err(error) => CallToolResult::error(vec![ContentBlock::text(error.to_string())]), }) }) } @@ -686,9 +651,8 @@ fn permission_granted(result: &ElicitResult) -> bool { fn permission_unsupported(tool_name: &str) -> String { format!( - "Permission required for {tool_name}, but the connected client cannot prompt the user \ - (MCP elicitation over protocol 2026-07-28 or newer). Use --permission-mode always-allow \ - or an elicitation-capable client." + "Permission required for {tool_name}, but {CLIENT_CANNOT_ELICIT}. \ + Use --permission-mode always-allow or an elicitation-capable client." ) } @@ -730,14 +694,6 @@ fn background_bash_args(request: &CallToolRequestParams) -> Option { args.run_in_background.is_some_and(|background| background).then_some(args) } -// TaskManager has no Drop of its own: aborting the tasks here drops their -// in-flight command futures, whose kill_on_drop reaps the child processes. -impl Drop for CodingMcp { - fn drop(&mut self) { - self.task_manager.shutdown(); - } -} - impl Debug for CodingMcp { fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { formatter @@ -748,12 +704,6 @@ impl Debug for CodingMcp { } } -impl Default for CodingMcp { - fn default() -> Self { - Self::new() - } -} - #[cfg(feature = "test-helpers")] impl CodingMcp { pub fn with_web_fetcher(mut self, web_fetcher: WebFetcher) -> Self { @@ -802,74 +752,69 @@ mod tests { #[test] fn args_default_permission_mode_is_always_allow() { - let args = CodingMcpArgs::from_args(vec![]).unwrap(); + let args = CodingMcpArgs::try_parse_from(Vec::<&str>::new()).unwrap(); assert_eq!(args.permission_mode, PermissionMode::AlwaysAllow); } #[test] fn args_default_lsp_enabled() { - let args = CodingMcpArgs::from_args(vec![]).unwrap(); + let args = CodingMcpArgs::try_parse_from(Vec::<&str>::new()).unwrap(); assert!(!args.disable_lsp); } #[test] fn args_parses_disable_lsp() { - let args = CodingMcpArgs::from_args(vec!["--disable-lsp".into()]).unwrap(); + let args = CodingMcpArgs::try_parse_from(["--disable-lsp"]).unwrap(); assert!(args.disable_lsp); } #[test] fn disabled_lsp_instructions_omit_lsp_tools() { - let instructions = CodingMcp::new().build_instructions(); + let instructions = CodingMcp::new("/workspace").build_instructions(); assert!(!instructions.contains("lsp_check_errors")); assert!(!instructions.contains("lsp_symbol")); } #[test] fn args_parses_always_allow() { - let args = CodingMcpArgs::from_args(vec!["--permission-mode".into(), "always-allow".into()]).unwrap(); + let args = CodingMcpArgs::try_parse_from(["--permission-mode", "always-allow"]).unwrap(); assert_eq!(args.permission_mode, PermissionMode::AlwaysAllow); } #[test] fn args_parses_auto() { - let args = CodingMcpArgs::from_args(vec!["--permission-mode".into(), "auto".into()]).unwrap(); + let args = CodingMcpArgs::try_parse_from(["--permission-mode", "auto"]).unwrap(); assert_eq!(args.permission_mode, PermissionMode::Auto); } #[test] fn args_parses_always_ask() { - let args = CodingMcpArgs::from_args(vec!["--permission-mode".into(), "always-ask".into()]).unwrap(); + let args = CodingMcpArgs::try_parse_from(["--permission-mode", "always-ask"]).unwrap(); assert_eq!(args.permission_mode, PermissionMode::AlwaysAsk); } #[test] fn args_rejects_invalid_permission_mode() { - assert!(CodingMcpArgs::from_args(vec!["--permission-mode".into(), "yolo".into()]).is_err()); + assert!(CodingMcpArgs::try_parse_from(["--permission-mode", "yolo"]).is_err()); } #[test] fn args_parses_repeated_rules_dirs() { - let args = CodingMcpArgs::from_args(vec![ - "--rules-dir".into(), - ".aether/skills".into(), - "--rules-dir".into(), - ".claude/rules".into(), - ]) - .unwrap(); + let args = + CodingMcpArgs::try_parse_from(["--rules-dir", ".aether/skills", "--rules-dir", ".claude/rules"]).unwrap(); assert_eq!(args.rules_dirs, vec![PathBuf::from(".aether/skills"), PathBuf::from(".claude/rules")]); } #[test] fn with_permission_mode_stores_mode() { - let mcp = CodingMcp::new().with_permission_mode(PermissionMode::AlwaysAsk); + let mcp = CodingMcp::new("/workspace").with_permission_mode(PermissionMode::AlwaysAsk); assert_eq!(mcp.permission_mode, PermissionMode::AlwaysAsk); } #[test] fn default_permission_mode_is_always_allow() { - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new("/workspace"); assert_eq!(mcp.permission_mode, PermissionMode::AlwaysAllow); } diff --git a/crates/mcp-servers/src/docs/coding_mcp.md b/crates/mcp-servers/src/docs/coding_mcp.md index a5fd819b2..fbb00da9b 100644 --- a/crates/mcp-servers/src/docs/coding_mcp.md +++ b/crates/mcp-servers/src/docs/coding_mcp.md @@ -7,21 +7,20 @@ Use the builder-style methods to configure the server: ```rust,ignore use mcp_servers::CodingMcp; -// Default: local filesystem tools, no LSP -let server = CodingMcp::new(); +// Default: local filesystem tools rooted at the project, no LSP +let server = CodingMcp::new("/my/project"); -// With LSP code intelligence and explicit read-rule directories -let server = CodingMcp::new() +// With LSP code intelligence for that root and explicit read-rule directories +let server = CodingMcp::new("/my/project") .with_rules_dirs(vec!["/my/project/.aether/skills".into(), "/my/project/.claude/rules".into()]) - .with_lsp("/my/project".into()) - .with_root_dir("/my/project".into()) + .with_lsp() .with_permission_mode(mcp_servers::PermissionMode::Auto); ``` For custom tool backends, use [`CodingMcp::with_tools`]: ```rust,ignore -let server = CodingMcp::with_tools(my_custom_tools); +let server = CodingMcp::with_tools("/my/project", my_custom_tools); ``` # Tools provided diff --git a/crates/mcp-servers/src/docs/mcp_builder_ext.md b/crates/mcp-servers/src/docs/mcp_builder_ext.md index f9e0051aa..ebaa3f94c 100644 --- a/crates/mcp-servers/src/docs/mcp_builder_ext.md +++ b/crates/mcp-servers/src/docs/mcp_builder_ext.md @@ -1,19 +1,18 @@ Extension trait that registers all built-in MCP server factories onto an [`McpBuilder`](aether_core::mcp::McpBuilder). -Call [`with_builtin_servers`](McpBuilderExt::with_builtin_servers) to register in-memory server factories for all built-in servers (coding, skills, subagents, review, tasks). Loading `mcp.json` records cloneable server specifications; the concrete servers are created only when [`McpBuilder::spawn`](aether_core::mcp::McpBuilder::spawn) runs. - -Install [`AgentDeps`](aether_core::core::AgentDeps) separately with [`with_agent_deps`](aether_core::mcp::McpBuilder::with_agent_deps). At spawn time every factory receives [`RuntimeServices`](aether_core::mcp::RuntimeServices), containing those dependencies, the builder root directory, the live [`McpHandle`](aether_core::mcp::McpHandle), and immutable shell environment entries for the session's deferred-tool gateway. Factory registration does not capture these runtime values. +Servers are constructed when [`McpBuilder::spawn`](aether_core::mcp::McpBuilder::spawn) runs. Invalid `args` for a server fail `spawn` with [`McpSpawnError::ServerFactory`](aether_core::mcp::McpSpawnError::ServerFactory) instead of falling back to defaults. # Usage ```rust,ignore use mcp_servers::McpBuilderExt; use aether_core::mcp::mcp; +use aether_core::mcp::McpConfig; let builder = mcp("/my/project") .with_agent_deps(deps) .with_builtin_servers() - .from_json_files(&["mcp.json"]) + .with_config(McpConfig::from_json_file("mcp.json").unwrap()) .unwrap(); ``` @@ -24,4 +23,3 @@ let builder = mcp("/my/project") - [`TasksMcp`](crate::TasksMcp) -- Task management - [`SubAgentsMcp`](crate::SubAgentsMcp) -- Sub-agent orchestration - [`ReviewMcp`](crate::ReviewMcp) -- Markdown artifact review workflow - diff --git a/crates/mcp-servers/src/docs/skills_mcp.md b/crates/mcp-servers/src/docs/skills_mcp.md index 943454c65..6da5fbb61 100644 --- a/crates/mcp-servers/src/docs/skills_mcp.md +++ b/crates/mcp-servers/src/docs/skills_mcp.md @@ -5,17 +5,16 @@ This server loads prompt artifacts from explicitly configured directories. It ex # Construction ```rust,ignore -use mcp_servers::SkillsMcp; +use clap::Parser; +use mcp_servers::{SkillsMcp, SkillsMcpArgs}; let server = SkillsMcp::new( + "/my/project", &["/my/project/.aether/skills".into(), "/my/project/.claude/rules".into()], ); -// From CLI args -let server = SkillsMcp::from_args(vec![ - "--dir".into(), "/my/project/.aether/skills".into(), - "--dir".into(), "/my/project/.claude/rules".into(), -]).unwrap(); +let args = SkillsMcpArgs::try_parse_from(["--dir", ".aether/skills", "--dir", ".claude/rules"])?; +let server = SkillsMcp::from_args(args, "/my/project".as_ref()); ``` # Tools provided diff --git a/crates/mcp-servers/src/docs/subagents_mcp.md b/crates/mcp-servers/src/docs/subagents_mcp.md index 902dfb0f3..32e2d2d36 100644 --- a/crates/mcp-servers/src/docs/subagents_mcp.md +++ b/crates/mcp-servers/src/docs/subagents_mcp.md @@ -5,14 +5,18 @@ Sub-agents are independent agent instances that run in parallel, each with their # Construction ```rust,ignore -use mcp_servers::SubAgentsMcp; +use clap::Parser; +use mcp_servers::{SubAgentsMcp, SubAgentsMcpArgs}; // Embedded: Aether supplies the complete runtime dependencies -let server = SubAgentsMcp::embedded_from_args(args, "/my/project".as_ref(), deps).unwrap(); +let server = SubAgentsMcp::new("/my/project".into(), deps); -// Standalone: discovers the registry from .aether/settings.json -let server = SubAgentsMcp::standalone("/my/project".into()).unwrap(); -let server = SubAgentsMcp::standalone_from_args(vec!["--project-root".into(), ".".into()]).unwrap(); +// Embedded, from CLI args resolved against a base directory +let args = SubAgentsMcpArgs::try_parse_from(["--project-root", "sub"])?; +let server = SubAgentsMcp::from_args(args.clone(), "/my/project".as_ref(), deps); + +// Standalone: discovers agents from the project's settings +let server = SubAgentsMcp::standalone(args, "/my/project".as_ref())?; ``` # Tools provided @@ -21,6 +25,6 @@ let server = SubAgentsMcp::standalone_from_args(vec!["--project-root".into(), ". # Agent catalog -Delegation targets come from the [`AgentRegistry`](aether_core::core::AgentRegistry) in the [`AgentDeps`](aether_core::core::AgentDeps) passed to [`SubAgentsMcp::embedded`](crate::SubAgentsMcp::embedded). Aether resolves the registry once at startup and threads it down through the deps, so SDK settings and other explicit settings sources stay authoritative at every delegation level — including sub-agents that themselves delegate. Embedded servers never reload agent definitions from the workspace. +Delegation targets come from the [`AgentRegistry`](aether_core::core::AgentRegistry) in the [`AgentDeps`](aether_core::core::AgentDeps) passed to [`SubAgentsMcp::new`](crate::SubAgentsMcp::new). Aether resolves the registry once at startup and threads it down through the deps, so SDK settings and other explicit settings sources stay authoritative at every delegation level — including sub-agents that themselves delegate. Embedded servers never reload agent definitions from the workspace. Standalone construction through [`SubAgentsMcp::standalone`](crate::SubAgentsMcp::standalone) or CLI arguments has no host to inherit from, so it discovers agents from `.aether/settings.json` in the project root and seeds its own deps with them. Each definition specifies a name, model, system prompt, and available tools. diff --git a/crates/mcp-servers/src/docs/tasks_mcp.md b/crates/mcp-servers/src/docs/tasks_mcp.md index 823321efc..d527b478a 100644 --- a/crates/mcp-servers/src/docs/tasks_mcp.md +++ b/crates/mcp-servers/src/docs/tasks_mcp.md @@ -10,16 +10,18 @@ Provides tools for creating, listing, and updating tasks organized into trees wi # Construction ```rust,ignore -use mcp_servers::TasksMcp; +use clap::Parser; +use mcp_servers::{TasksMcp, TasksMcpArgs}; // Session-scoped let server = TasksMcp::new(); -// Persistent -let server = TasksMcp::new_persistent("/my/project/.aether-tasks".into()); +// Persistent, in /my/project/.aether-tasks +let server = TasksMcp::new_persistent("/my/project"); -// From CLI args (e.g. --dir /path) -let server = TasksMcp::from_args(vec!["--dir".into(), ".".into()]).unwrap(); +// From CLI args (e.g. --dir /path), with a relative --dir resolved against a base directory +let args = TasksMcpArgs::try_parse_from(["--dir", "."])?; +let server = TasksMcp::from_args(args, "/my/project".as_ref()); ``` # Tools provided diff --git a/crates/mcp-servers/src/docs/tasks_module.md b/crates/mcp-servers/src/docs/tasks_module.md index 970c7613c..a4b75bb41 100644 --- a/crates/mcp-servers/src/docs/tasks_module.md +++ b/crates/mcp-servers/src/docs/tasks_module.md @@ -19,11 +19,11 @@ Provides a standalone MCP server for hierarchical task tracking with support for ```rust,ignore use mcp_servers::TasksMcp; -use mcp_utils::ServiceExt; +use mcp_utils::server::McpServer; // Session-scoped (temp dir, auto-cleanup) -let server = TasksMcp::new().into_dyn(); +let server = McpServer::new(TasksMcp::new()); // Persistent across sessions -let server = TasksMcp::new_persistent(".aether-tasks".into()).into_dyn(); +let server = McpServer::new(TasksMcp::new_persistent(".")); ``` diff --git a/crates/mcp-servers/src/error.rs b/crates/mcp-servers/src/error.rs deleted file mode 100644 index 6c26c9b82..000000000 --- a/crates/mcp-servers/src/error.rs +++ /dev/null @@ -1,9 +0,0 @@ -use thiserror::Error; - -#[derive(Debug, Error)] -pub enum ServerInitError { - #[error("Failed to parse arguments: {0}")] - InvalidArgs(#[from] clap::Error), - #[error("{0}")] - Other(String), -} diff --git a/crates/mcp-servers/src/lib.rs b/crates/mcp-servers/src/lib.rs index 1b332599a..1a5ea61f3 100644 --- a/crates/mcp-servers/src/lib.rs +++ b/crates/mcp-servers/src/lib.rs @@ -1,16 +1,20 @@ #![doc = include_str!("../README.md")] -pub mod error; +#[cfg(any(feature = "coding", feature = "subagents"))] +mod background_tasks; + +#[cfg(any(feature = "coding", feature = "review"))] +mod request_context; pub mod workspace_paths; #[cfg(any(test, feature = "test-helpers"))] pub mod testing; -#[cfg(any(feature = "all", feature = "subagents"))] +#[cfg(feature = "subagents")] pub mod setup; -#[cfg(any(feature = "all", feature = "subagents"))] -pub use setup::McpBuilderExt; +#[cfg(feature = "subagents")] +pub use setup::{Builtin, McpBuilderExt, agent_mcp}; #[cfg(feature = "file-ops")] pub mod file_ops; @@ -33,7 +37,6 @@ pub mod subagents; #[cfg(feature = "review")] pub mod review; -// Re-export primary types for convenience #[cfg(feature = "coding")] pub use coding::{CodingMcp, CodingMcpArgs, CodingTools, DefaultCodingTools, PermissionMode}; @@ -53,3 +56,14 @@ pub use subagents::{SubAgentsMcp, SubAgentsMcpArgs}; pub use review::{ DocumentSource, HtmlAnnotation, HtmlSource, ReviewArtifactInput, ReviewArtifactOutput, ReviewMcp, ReviewMcpArgs, }; + +#[cfg(any(feature = "coding", feature = "skills", feature = "tasks", feature = "review"))] +pub(crate) fn server_config( + name: &str, + capabilities: rmcp::model::ServerCapabilities, + instructions: impl Into, +) -> rmcp::model::ServerConfig { + rmcp::model::ServerConfig::new(capabilities) + .with_server_info(rmcp::model::Implementation::new(name, env!("CARGO_PKG_VERSION"))) + .with_instructions(instructions) +} diff --git a/crates/mcp-servers/src/request_context.rs b/crates/mcp-servers/src/request_context.rs new file mode 100644 index 000000000..717f697ce --- /dev/null +++ b/crates/mcp-servers/src/request_context.rs @@ -0,0 +1,37 @@ +use rmcp::RoleServer; +use rmcp::model::{ErrorData, InputResponses}; +use rmcp::service::{ElicitationMode, RequestContext}; +use serde::de::DeserializeOwned; + +pub(crate) const CLIENT_CANNOT_ELICIT: &str = + "the connected client cannot prompt the user (MCP elicitation over protocol 2026-07-28 or newer)"; + +pub(crate) fn supports_elicitation(context: &RequestContext, mode: ElicitationMode) -> bool { + let Some(elicitation) = context.client_capabilities().and_then(|capabilities| capabilities.elicitation) else { + return false; + }; + match mode { + ElicitationMode::Form => elicitation.form.is_some() || elicitation.url.is_none(), + ElicitationMode::Url => elicitation.url.is_some(), + _ => false, + } +} + +#[cfg(feature = "coding")] +pub(crate) fn require_tasks(context: &RequestContext) -> Result<(), ErrorData> { + if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks()) { + Ok(()) + } else { + let tasks = rmcp::model::ClientCapabilities::builder().enable_tasks().build(); + Err(ErrorData::missing_required_client_capability(tasks)) + } +} + +pub(crate) fn parse_response(responses: &InputResponses, key: &str) -> Result { + let response = responses + .get(key) + .ok_or_else(|| ErrorData::invalid_params(format!("missing input response for '{key}'"), None))?; + + serde_json::from_value(response.clone()) + .map_err(|e| ErrorData::invalid_params(format!("invalid input response for '{key}': {e}"), None)) +} diff --git a/crates/mcp-servers/src/review/server.rs b/crates/mcp-servers/src/review/server.rs index 65e03ddf6..4c771c07f 100644 --- a/crates/mcp-servers/src/review/server.rs +++ b/crates/mcp-servers/src/review/server.rs @@ -1,6 +1,6 @@ use super::tools::{ReviewArtifactInput, ReviewArtifactOutput, ReviewArtifactTool}; -use crate::error::ServerInitError; -use crate::workspace_paths::{current_dir, resolve_path}; +use crate::server_config; +use crate::workspace_paths::resolve_root; use clap::Parser; use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, @@ -9,29 +9,21 @@ use rmcp::{ tool::{InputResponses, RequestState, schema_for_output}, wrapper::Parameters, }, - model::{CallToolResponse, Implementation, ServerCapabilities, ServerConfig}, + model::{CallToolResponse, ServerCapabilities, ServerConfig}, service::RequestContext, tool, tool_handler, tool_router, }; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; #[derive(Debug, Clone, Parser)] -#[command(name = "review-mcp")] +#[command(name = "review", no_binary_name = true)] pub struct ReviewMcpArgs { /// Workspace root used to resolve relative artifact paths. #[arg(long)] pub root_dir: Option, } -impl ReviewMcpArgs { - pub fn from_args(args: Vec) -> Result { - let mut full_args = vec!["review-mcp".to_string()]; - full_args.extend(args); - Self::try_parse_from(full_args).map_err(ServerInitError::InvalidArgs) - } -} - #[doc = include_str!("../docs/review_mcp.md")] #[derive(Clone)] pub struct ReviewMcp { @@ -39,23 +31,18 @@ pub struct ReviewMcp { review_artifact: Arc, } -#[tool_router] impl ReviewMcp { - pub fn new() -> Self { - Self::at_root(current_dir()) - } - - pub fn from_args(args: Vec) -> Result { - let parsed = ReviewMcpArgs::from_args(args)?; - Ok(Self::at_root(parsed.root_dir.unwrap_or_else(current_dir))) + pub fn new(root_dir: PathBuf) -> Self { + Self { tool_router: Self::tool_router(), review_artifact: Arc::new(ReviewArtifactTool::new(root_dir)) } } - pub fn from_args_with_base_dir(args: Vec, base_dir: &std::path::Path) -> Result { - let parsed = ReviewMcpArgs::from_args(args)?; - let root_dir = parsed.root_dir.map_or_else(|| base_dir.to_path_buf(), |path| resolve_path(base_dir, path)); - Ok(Self::at_root(root_dir)) + pub fn from_args(args: ReviewMcpArgs, base_dir: &Path) -> Self { + Self::new(resolve_root(base_dir, args.root_dir)) } +} +#[tool_router] +impl ReviewMcp { #[doc = include_str!("tools/review_artifact/description.md")] #[tool( annotations(read_only_hint = true, destructive_hint = false, idempotent_hint = false, open_world_hint = false), @@ -76,20 +63,10 @@ impl ReviewMcp { #[tool_handler(router = self.tool_router)] impl ServerHandler for ReviewMcp { fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("review-mcp", "0.1.0")) - .with_instructions(include_str!("./instructions.md")) - } -} - -impl ReviewMcp { - fn at_root(root_dir: PathBuf) -> Self { - Self { tool_router: Self::tool_router(), review_artifact: Arc::new(ReviewArtifactTool::new(root_dir)) } - } -} - -impl Default for ReviewMcp { - fn default() -> Self { - Self::new() + server_config( + "review-mcp", + ServerCapabilities::builder().enable_tools().build(), + include_str!("./instructions.md"), + ) } } diff --git a/crates/mcp-servers/src/review/tools/review_artifact/mod.rs b/crates/mcp-servers/src/review/tools/review_artifact/mod.rs index fe0b9d2e0..6b2bff283 100644 --- a/crates/mcp-servers/src/review/tools/review_artifact/mod.rs +++ b/crates/mcp-servers/src/review/tools/review_artifact/mod.rs @@ -1,6 +1,6 @@ +use crate::request_context::{CLIENT_CANNOT_ELICIT, parse_response, supports_elicitation}; use crate::review::html_review::{Artifact, PendingReviews, ReviewServer}; use crate::workspace_paths::resolve_path; -use mcp_utils::server::mrtr::{ELICITATION_UNSUPPORTED, ElicitationMode, elicitation_supported, parse_response}; use reqwest::Url; use rmcp::{ ErrorData as McpError, RoleServer, @@ -13,7 +13,7 @@ use rmcp::{ ElicitationAction, ElicitationSchema, EnumSchema, InputRequest, InputRequests, InputRequiredResult, RequestMetaObject, }, - service::RequestContext, + service::{ElicitationMode, RequestContext}, }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -110,8 +110,8 @@ impl ReviewArtifactTool { ReviewArtifactInput::Markdown { .. } => ElicitationMode::Form, ReviewArtifactInput::Html { .. } => ElicitationMode::Url, }; - if !elicitation_supported(context.client_capabilities().as_ref(), mode) { - return Err(ELICITATION_UNSUPPORTED.to_string()); + if !supports_elicitation(context, mode) { + return Err(format!("Reviewing an artifact requires user input, but {CLIENT_CANNOT_ELICIT}.")); } let (params, review_token) = match input { diff --git a/crates/mcp-servers/src/setup.rs b/crates/mcp-servers/src/setup.rs index df3bbcdba..b4de43ae8 100644 --- a/crates/mcp-servers/src/setup.rs +++ b/crates/mcp-servers/src/setup.rs @@ -1,127 +1,69 @@ use crate::coding::tools::bash::BashEnvironment; -use crate::workspace_paths::resolve_path; -use crate::{CodingMcp, CodingMcpArgs, DefaultCodingTools, ReviewMcp, SkillsMcp, SubAgentsMcp, TasksMcp}; -use aether_core::mcp::{McpBuilder, RuntimeServices}; -use futures::FutureExt; -use mcp_utils::ServiceExt; -use tracing::{debug, warn}; - -pub fn progressive_discovery_instructions() -> String { - include_str!("progressive_discovery_instructions.md").to_string() -} +use crate::{CodingMcp, DefaultCodingTools, ReviewMcp, SkillsMcp, SubAgentsMcp, TasksMcp}; +use aether_core::core::AgentDeps; +use aether_core::mcp::{AETHER_MCP_IPC_SOCKET, McpBuilder, RuntimeServices, ToolFilter, mcp}; +use clap::{Parser, ValueEnum}; +use mcp_utils::server::McpServer; +use std::path::Path; #[doc = include_str!("docs/mcp_builder_ext.md")] pub trait McpBuilderExt { - /// Registers built-in in-memory MCP factories. Servers are constructed at - /// spawn time with the runtime's root directory, dependencies, and MCP handle. fn with_builtin_servers(self) -> Self; } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum Builtin { + Coding, + Skills, + Tasks, + Subagents, + Review, +} + +pub fn agent_mcp(root: impl AsRef, deps: AgentDeps, tools: ToolFilter) -> McpBuilder { + mcp(root).with_agent_deps(deps).with_tool_filter(tools).with_builtin_servers() +} + impl McpBuilderExt for McpBuilder { fn with_builtin_servers(self) -> Self { - self.register_in_memory_server( - "coding", - Box::new(|spec, services| { - async move { - let parsed = match CodingMcpArgs::from_args(spec.args) { - Ok(args) => args, - Err(e) => { - warn!("CodingMcp args parse failed: {e}, using defaults"); - CodingMcpArgs::default() - } - }; - let CodingMcpArgs { permission_mode, mut rules_dirs, disable_lsp, root_dir: arg_root_dir } = parsed; - let root_dir = arg_root_dir - .map_or_else(|| services.root_dir.clone(), |path| resolve_path(&services.root_dir, path)); - rules_dirs = rules_dirs.into_iter().map(|path| resolve_path(&root_dir, path)).collect(); - debug!( - "CodingMcp created, disable_lsp={}, permission_mode={:?}, rules_dirs={}", - disable_lsp, - permission_mode, - rules_dirs.len() - ); - let tools = DefaultCodingTools::new().with_bash_environment(coding_bash_environment(&services)); - let server = CodingMcp::with_tools(tools) - .with_rules_dirs(rules_dirs) - .with_root_dir(root_dir.clone()) - .with_permission_mode(permission_mode); - let server = if disable_lsp { server } else { server.with_lsp(root_dir) }; - server.into_dyn() - } - .boxed() - }), - ) - .register_in_memory_server( - "skills", - Box::new(|spec, services| { - async move { - SkillsMcp::from_args_with_base_dir(spec.args, &services.root_dir) - .map_err(|error| { - warn!("Failed to parse SkillsMcp args: {error}, using defaults"); - error - }) - .unwrap_or_else(|_| SkillsMcp::new(&[]).with_root_dir(services.root_dir.clone())) - .into_dyn() - } - .boxed() - }), - ) - .register_in_memory_server( - "subagents", - Box::new(|spec, services| { - async move { - let agent_deps = services.agent_deps.clone(); - SubAgentsMcp::embedded_from_args(spec.args, &services.root_dir, agent_deps.clone()) - .map_err(|error| { - warn!("Failed to parse SubAgentsMcp args: {error}, using defaults"); - error - }) - .unwrap_or_else(|_| SubAgentsMcp::embedded(services.root_dir.clone(), agent_deps)) - .into_dyn() - } - .boxed() - }), - ) - .register_in_memory_server( - "review", - Box::new(|spec, services| { - async move { - ReviewMcp::from_args_with_base_dir(spec.args, &services.root_dir) - .inspect_err(|error| { - warn!("Failed to parse ReviewMcp args: {error}, using workspace root"); - }) - .unwrap_or_else(|_| { - ReviewMcp::from_args_with_base_dir(Vec::new(), &services.root_dir) - .expect("empty review MCP arguments are valid") - }) - .into_dyn() - } - .boxed() - }), - ) - .register_in_memory_server( - "tasks", - Box::new(|spec, services| { - async move { - TasksMcp::from_args_with_base_dir(spec.args, &services.root_dir) - .unwrap_or_else(|e| { - tracing::warn!("Failed to parse TasksMcp args: {e}, using defaults"); - TasksMcp::new() - }) - .into_dyn() - } - .boxed() - }), - ) - .with_progressive_discovery_instructions(progressive_discovery_instructions()) + Builtin::value_variants().iter().fold(self, |builder, &builtin| { + builder.register_in_memory_server(builtin.name(), move |args, services| Ok(builtin.build(args, services)?)) + }) } } -fn coding_bash_environment(services: &RuntimeServices) -> BashEnvironment { - services - .shell_environment - .iter() - .fold(BashEnvironment::new().with_current_exe_dir_on_path(), |environment, (name, value)| { - environment.with_var(name, value) +impl Builtin { + pub fn name(self) -> &'static str { + match self { + Self::Coding => "coding", + Self::Skills => "skills", + Self::Tasks => "tasks", + Self::Subagents => "subagents", + Self::Review => "review", + } + } + + pub fn build(self, args: Vec, services: &RuntimeServices) -> Result { + let root = services.root_dir.as_path(); + Ok(match self { + Self::Coding => CodingMcp::from_args(parse(args)?, root, coding_tools(services)).into(), + Self::Skills => SkillsMcp::from_args(parse(args)?, root).into(), + Self::Tasks => TasksMcp::from_args(parse(args)?, root).into(), + Self::Review => ReviewMcp::from_args(parse(args)?, root).into(), + Self::Subagents => SubAgentsMcp::from_args(parse(args)?, root, services.agent_deps.clone()).into(), }) + } +} + +fn parse(args: Vec) -> Result { + T::try_parse_from(args) +} + +fn coding_tools(services: &RuntimeServices) -> DefaultCodingTools { + let environment = BashEnvironment::new().with_current_exe_dir_on_path(); + let environment = match &services.deferred_tools_socket { + Some(socket) => environment.with_var(AETHER_MCP_IPC_SOCKET, socket.to_string_lossy()), + None => environment, + }; + DefaultCodingTools::new().with_bash_environment(environment) } diff --git a/crates/mcp-servers/src/skills/server.rs b/crates/mcp-servers/src/skills/server.rs index 15c3ce168..fe5f7f246 100644 --- a/crates/mcp-servers/src/skills/server.rs +++ b/crates/mcp-servers/src/skills/server.rs @@ -1,3 +1,4 @@ +use crate::server_config; use clap::Parser; use rmcp::{ ErrorData as McpError, RoleServer, ServerHandler, @@ -6,8 +7,8 @@ use rmcp::{ wrapper::{Json, Parameters}, }, model::{ - GetPromptRequestParams, GetPromptResponse, Implementation, ListPromptsResult, PaginatedRequestParams, Prompt, - PromptArgument, PromptMessage, Role, ServerCapabilities, ServerConfig, + GetPromptRequestParams, GetPromptResponse, ListPromptsResult, PaginatedRequestParams, Prompt, PromptArgument, + PromptMessage, Role, ServerCapabilities, ServerConfig, }, service::RequestContext, tool, tool_handler, tool_router, @@ -23,15 +24,13 @@ use utils::substitution::substitute_parameters; use super::tools::{ ListSkillsInput, ListSkillsOutput, LoadSkillsInput, LoadSkillsOutput, SkillFile, SkillListItem, SkillRequest, }; -use crate::{ - error::ServerInitError, - workspace_paths::{WorkspacePaths, resolve_path}, -}; +use crate::workspace_paths::{WorkspacePaths, resolve_path}; use aether_project::{PromptCatalog, PromptFile, SKILL_FILENAME}; use utils::display_meta::ToolDisplayMeta; /// CLI arguments for `SkillsMcp` server #[derive(Debug, Clone, Parser)] +#[command(name = "skills", no_binary_name = true)] pub struct SkillsMcpArgs { /// Prompt directories to scan directly. /// Can be specified multiple times: --dir .aether/skills --dir .claude/rules @@ -44,15 +43,6 @@ pub struct SkillsMcpArgs { pub notes_dir: Option, } -impl SkillsMcpArgs { - pub fn from_args(args: Vec) -> Result { - let mut full_args = vec!["skills-mcp".to_string()]; - full_args.extend(args); - - Self::try_parse_from(full_args).map_err(ServerInitError::InvalidArgs) - } -} - #[doc = include_str!("../docs/skills_mcp.md")] #[derive(Clone)] pub struct SkillsMcp { @@ -86,30 +76,15 @@ enum SkillFileError { } impl SkillsMcp { - pub fn new(prompt_dirs: &[PathBuf]) -> Self { + pub fn new(root_dir: impl Into, prompt_dirs: &[PathBuf]) -> Self { let catalog = PromptCatalog::from_dirs(prompt_dirs); - Self { - catalog: Arc::new(catalog), - tool_router: Self::tool_router(), - root_dir: crate::workspace_paths::current_dir(), - } - } - - pub fn from_args(args: Vec) -> Result { - let parsed_args = SkillsMcpArgs::from_args(args)?; - Ok(Self::new(&parsed_args.dirs)) - } - - pub fn from_args_with_base_dir(args: Vec, base_dir: &Path) -> Result { - let parsed_args = SkillsMcpArgs::from_args(args)?; - let dirs = parsed_args.dirs.into_iter().map(|path| resolve_path(base_dir, path)).collect::>(); - Ok(Self::new(&dirs).with_root_dir(base_dir.to_path_buf())) + Self { catalog: Arc::new(catalog), tool_router: Self::tool_router(), root_dir: root_dir.into() } } - pub fn with_root_dir(mut self, root_dir: PathBuf) -> Self { - self.root_dir = root_dir; - self + pub fn from_args(args: SkillsMcpArgs, base_dir: &Path) -> Self { + let dirs = args.dirs.into_iter().map(|path| resolve_path(base_dir, path)).collect::>(); + Self::new(base_dir, &dirs) } fn is_directory_prompt(prompt: &PromptFile) -> bool { @@ -256,9 +231,11 @@ impl SkillsMcp { #[tool_handler(router = self.tool_router)] impl ServerHandler for SkillsMcp { fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_prompts().enable_tools().build()) - .with_server_info(Implementation::new("skills-mcp", "0.1.0")) - .with_instructions(include_str!("./instructions.md")) + server_config( + "skills-mcp", + ServerCapabilities::builder().enable_prompts().enable_tools().build(), + include_str!("./instructions.md"), + ) } async fn list_prompts( @@ -384,45 +361,31 @@ mod tests { #[test] fn skills_args_require_at_least_one_dir() { - let parsed = SkillsMcpArgs::from_args(vec!["--notes-dir".into(), ".aether/notes".into()]); + let parsed = SkillsMcpArgs::try_parse_from(["--notes-dir", ".aether/notes"]); assert!(parsed.is_err()); } #[test] fn skills_args_parse_repeated_dirs() { - let parsed = SkillsMcpArgs::from_args(vec![ - "--dir".into(), - ".aether/skills".into(), - "--dir".into(), - ".claude/rules".into(), - ]) - .unwrap(); + let parsed = SkillsMcpArgs::try_parse_from(["--dir", ".aether/skills", "--dir", ".claude/rules"]).unwrap(); assert_eq!(parsed.dirs, vec![PathBuf::from(".aether/skills"), PathBuf::from(".claude/rules")]); } #[test] fn skills_args_accepts_deprecated_notes_dir() { - let parsed = SkillsMcpArgs::from_args(vec![ - "--dir".into(), - ".aether/skills".into(), - "--notes-dir".into(), - ".aether/notes".into(), - ]) - .unwrap(); + let parsed = + SkillsMcpArgs::try_parse_from(["--dir", ".aether/skills", "--notes-dir", ".aether/notes"]).unwrap(); assert_eq!(parsed.dirs, vec![PathBuf::from(".aether/skills")]); assert_eq!(parsed.notes_dir, Some(PathBuf::from(".aether/notes"))); } #[test] - fn from_args_with_base_dir_resolves_relative_dirs() { + fn from_args_roots_server_at_base_dir() { let workspace = PathBuf::from("/workspace"); - let server = SkillsMcp::from_args_with_base_dir( - vec!["--dir".into(), ".aether/skills".into(), "--notes-dir".into(), ".aether/notes".into()], - &workspace, - ) - .unwrap(); + let args = SkillsMcpArgs::try_parse_from(["--dir", ".aether/skills"]).unwrap(); + let server = SkillsMcp::from_args(args, &workspace); assert_eq!(server.root_dir, workspace); } @@ -453,7 +416,7 @@ mod tests { "---\ndescription: Rule only\nagent-invocable: false\ntriggers:\n read:\n - \"**/*.rs\"\n---\n# Rule", ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let Json(output) = server.list_skills(Parameters(ListSkillsInput::default())).await; assert_eq!(output.status, "success"); @@ -492,7 +455,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let info = ServerHandler::get_info(&server); let instructions = info.instructions.expect("skills server should provide instructions"); @@ -512,7 +475,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "user-only", None).await; assert!(result.content.is_none()); @@ -528,7 +491,7 @@ mod tests { "---\ndescription: Rule-only prompt\nagent-invocable: false\ntriggers:\n read:\n - \"**/*.rs\"\n---\n# Rule", ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "rule-only", None).await; assert!(result.content.is_none()); @@ -545,7 +508,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", None).await; assert_eq!(result.name, "test-skill"); @@ -565,7 +528,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "custom-name", None).await; assert_eq!(result.name, "custom-name"); @@ -584,7 +547,7 @@ mod tests { "---\nname: rust-rules\ndescription: Rust rules\nagent-invocable: true\ntriggers:\n read:\n - \"**/*.rs\"\n---\nUse Rust conventions.", ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "rust-rules", None).await; assert_eq!(result.name, "rust-rules"); @@ -603,7 +566,7 @@ mod tests { "---\ndescription: Rust rules\nagent-invocable: true\ntriggers:\n read:\n - \"**/*.rs\"\n---\nUse Rust conventions.", ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "rust-rules", None).await; assert_eq!(result.path, "rust-rules.md"); @@ -622,7 +585,7 @@ mod tests { "---\ndescription: Rust rules\nagent-invocable: true\ntriggers:\n read:\n - \"**/*.rs\"\n---\nUse Rust conventions.", ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "rust-rules", Some("details.md")).await; assert!(result.content.is_none()); @@ -640,7 +603,7 @@ mod tests { &[("traits.md", "# Traits content")], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some("traits.md")).await; assert_eq!(result.path, "traits.md"); @@ -658,7 +621,7 @@ mod tests { &[("references/REF.md", "# Reference")], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some("references/REF.md")).await; assert_eq!(result.path, "references/REF.md"); @@ -670,7 +633,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); create_skill(temp_dir.path(), "test-skill", "---\ndescription: Test\nagent-invocable: true\n---\n# Test", &[]); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some("/etc/passwd")).await; assert!(result.error.unwrap().contains("Absolute paths")); @@ -681,7 +644,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); create_skill(temp_dir.path(), "test-skill", "---\ndescription: Test\nagent-invocable: true\n---\n# Test", &[]); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some("../other-skill/SKILL.md")).await; assert!(result.error.unwrap().contains("traversal")); @@ -692,7 +655,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); create_skill(temp_dir.path(), "test-skill", "---\ndescription: Test\nagent-invocable: true\n---\n# Test", &[]); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some(".")).await; assert!(result.error.unwrap().contains("directory")); @@ -713,7 +676,7 @@ mod tests { ], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", None).await; assert_eq!(result.available_files.len(), 3); @@ -728,7 +691,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); fs::create_dir_all(temp_dir.path().join("skills")).unwrap(); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "nonexistent", None).await; assert!(result.error.unwrap().contains("not found")); @@ -739,7 +702,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); create_skill(temp_dir.path(), "test-skill", "---\ndescription: Test\nagent-invocable: true\n---\n# Test", &[]); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", Some("nonexistent.md")).await; assert!(result.error.unwrap().contains("not found")); @@ -761,7 +724,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let expander = ShellExpander::new(); let requests = vec![ SkillRequest { name: "rust".to_string(), path: None }, @@ -796,7 +759,7 @@ mod tests { let temp_dir = TempDir::new().unwrap(); create_skill(temp_dir.path(), "exists", "---\ndescription: Exists\nagent-invocable: true\n---\n# Exists", &[]); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let expander = ShellExpander::new(); let requests = vec![ SkillRequest { name: "exists".to_string(), path: None }, @@ -828,7 +791,7 @@ mod tests { create_skill(dir_a.path(), "rust", "---\ndescription: Rust A\nagent-invocable: true\n---\n# From A", &[]); create_skill(dir_b.path(), "rust", "---\ndescription: Rust B\nagent-invocable: true\n---\n# From B", &[]); - let server = SkillsMcp::new(&[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); + let server = SkillsMcp::new(dir_a.path(), &[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); let result = load(&server, "rust", None).await; assert!(result.content.as_ref().unwrap().contains("From B")); @@ -841,7 +804,7 @@ mod tests { create_skill(dir_a.path(), "rust", "---\ndescription: Rust\nagent-invocable: true\n---\n# Rust", &[]); create_skill(dir_b.path(), "python", "---\ndescription: Python\nagent-invocable: true\n---\n# Python", &[]); - let server = SkillsMcp::new(&[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); + let server = SkillsMcp::new(dir_a.path(), &[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); let rust = load(&server, "rust", None).await; assert!(rust.content.is_some()); @@ -858,7 +821,7 @@ mod tests { let dir_b = TempDir::new().unwrap(); create_skill(dir_b.path(), "rust", "---\ndescription: Rust\nagent-invocable: true\n---\n# Rust", &[]); - let server = SkillsMcp::new(&[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); + let server = SkillsMcp::new(dir_a.path(), &[dir_a.path().to_path_buf(), dir_b.path().to_path_buf()]); let result = load(&server, "rust", None).await; assert!(result.content.is_some()); @@ -875,7 +838,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", None).await; let content = result.content.expect("content should be present"); @@ -894,7 +857,7 @@ mod tests { &[], ); - let server = SkillsMcp::new(&[temp_dir.path().to_path_buf()]); + let server = SkillsMcp::new(temp_dir.path(), &[temp_dir.path().to_path_buf()]); let result = load(&server, "test-skill", None).await; let content = result.content.expect("content should be present"); diff --git a/crates/mcp-servers/src/subagents/server.rs b/crates/mcp-servers/src/subagents/server.rs index 1e54271a1..072c252e6 100644 --- a/crates/mcp-servers/src/subagents/server.rs +++ b/crates/mcp-servers/src/subagents/server.rs @@ -1,80 +1,58 @@ +use crate::server_config; use aether_core::core::AgentDeps; use aether_core::events::{McpRequestInstrumentation, TraceContext}; -use aether_project::{AetherSettings, AgentCatalog}; +use aether_project::{AetherSettings, AgentCatalog, SettingsError}; use clap::Parser; -use mcp_utils::server::tasks::{BACKGROUND_TASK_TTL_MS, require_tasks_capability}; use rmcp::model::{CallToolResult, Task}; use rmcp::{ ErrorData, RoleServer, ServerHandler, handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{ - CallToolResponse, CancelTaskParams, CreateTaskResult, GetTaskParams, GetTaskResult, Implementation, - ServerCapabilities, ServerConfig, UpdateTaskParams, - }, + model::{CallToolResponse, CreateTaskResult, ServerCapabilities, ServerConfig}, service::RequestContext, - task_manager::{TaskContext, TaskExit, TaskManager, TaskOptions}, + task_manager::TaskExit, tool, tool_handler, tool_router, }; use std::path::{Path, PathBuf}; use super::progress::SubAgentProgressSink; use super::tools::{AgentExecutor, SpawnSubAgentsInput, SpawnSubAgentsOutput}; -use crate::error::ServerInitError; -use crate::workspace_paths::resolve_path; +use crate::background_tasks::{BackgroundTasks, delegate_background_tasks}; +use crate::request_context::require_tasks; +use crate::workspace_paths::resolve_root; const SPAWN_SUBAGENT_TOOL: &str = "spawn_subagent"; #[derive(Debug, Clone, Parser)] +#[command(name = "subagents", no_binary_name = true)] pub struct SubAgentsMcpArgs { /// Project root containing optional .aether/settings.json #[arg(long = "project-root", alias = "dir")] pub project_root: Option, } -impl SubAgentsMcpArgs { - pub fn from_args(args: Vec) -> Result { - let mut full_args = vec!["subagents-mcp".to_string()]; - full_args.extend(args); - Self::try_parse_from(full_args).map_err(ServerInitError::InvalidArgs) - } - - fn resolve_root(self, default_root: &Path) -> PathBuf { - self.project_root.map_or_else(|| default_root.to_path_buf(), |path| resolve_path(default_root, path)) - } -} - #[doc = include_str!("../docs/subagents_mcp.md")] pub struct SubAgentsMcp { tool_router: ToolRouter, - task_manager: TaskManager, + tasks: BackgroundTasks, project_root: PathBuf, agent_deps: AgentDeps, } impl SubAgentsMcp { - pub fn embedded(project_root: PathBuf, agent_deps: AgentDeps) -> Self { - Self { tool_router: Self::tool_router(), task_manager: TaskManager::new(), project_root, agent_deps } + pub fn new(project_root: PathBuf, agent_deps: AgentDeps) -> Self { + Self { tool_router: Self::tool_router(), tasks: BackgroundTasks::new(), project_root, agent_deps } } - pub fn embedded_from_args( - args: Vec, - base_dir: &Path, - agent_deps: AgentDeps, - ) -> Result { - Ok(Self::embedded(SubAgentsMcpArgs::from_args(args)?.resolve_root(base_dir), agent_deps)) + pub fn from_args(args: SubAgentsMcpArgs, base_dir: &Path, agent_deps: AgentDeps) -> Self { + Self::new(resolve_root(base_dir, args.project_root), agent_deps) } - pub fn standalone(project_root: PathBuf) -> Result { - let settings = AetherSettings::load_default(&project_root) - .map_err(|e| ServerInitError::Other(format!("Failed to load agents: {e}")))?; - let catalog = AgentCatalog::from_settings_or_empty(&project_root, settings) - .map_err(|e| ServerInitError::Other(format!("Failed to load agents: {e}")))?; + pub fn standalone(args: SubAgentsMcpArgs, base_dir: &Path) -> Result { + let project_root = resolve_root(base_dir, args.project_root); + let settings = AetherSettings::load_default(&project_root)?; + let catalog = AgentCatalog::from_settings_or_empty(&project_root, settings)?; let registry = catalog.registry().clone(); - Ok(Self::embedded(project_root, AgentDeps::default().with_agent_registry(registry))) - } - - pub fn standalone_from_args(args: Vec) -> Result { - Self::standalone(SubAgentsMcpArgs::from_args(args)?.resolve_root(Path::new("."))) + Ok(Self::new(project_root, AgentDeps::default().with_agent_registry(registry))) } fn build_instructions(&self) -> String { @@ -121,7 +99,7 @@ impl SubAgentsMcp { )); } if args.run_in_background { - require_tasks_capability(context)?; + require_tasks(context)?; } Ok(()) } @@ -148,19 +126,11 @@ impl SubAgentsMcp { tasks => format!("Running {} sub-agents", tasks.len()), }; let executor = self.executor(context, instrumentation.trace_context()); - let options = TaskOptions::new().with_ttl_ms(BACKGROUND_TASK_TTL_MS).with_status_message(status); - self.task_manager.spawn(options, move |task_context: TaskContext| { - Box::pin(async move { - tokio::select! { - () = task_context.cancelled() => Err(TaskExit::Cancelled), - output = executor.execute_tasks(args.tasks) => { - let result = output_result(output); - instrumentation.finish(result.as_ref().err().map(ToString::to_string).as_deref()); - result.map_err(TaskExit::Error) - } - } - }) + self.tasks.spawn(status, async move { + let result = output_result(executor.execute_tasks(args.tasks).await); + instrumentation.finish(result.as_ref().err().map(ToString::to_string).as_deref()); + result.map_err(TaskExit::Error) }) } } @@ -169,34 +139,14 @@ impl SubAgentsMcp { #[tool_handler(router = self.tool_router)] impl ServerHandler for SubAgentsMcp { fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().enable_tasks().build()) - .with_server_info(Implementation::new("subagents-mcp", "0.1.0")) - .with_instructions(self.build_instructions()) - } - - fn get_task( - &self, - request: GetTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.get_task(&request.task_id).map(GetTaskResult::new)) + server_config( + "subagents-mcp", + ServerCapabilities::builder().enable_tools().enable_tasks().build(), + self.build_instructions(), + ) } - fn update_task( - &self, - request: UpdateTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.update_task(&request.task_id, request.input_responses)) - } - - fn cancel_task( - &self, - request: CancelTaskParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(self.task_manager.cancel_task(&request.task_id)) - } + delegate_background_tasks!(tasks); } #[tool_router] @@ -232,12 +182,6 @@ impl SubAgentsMcp { } } -impl Drop for SubAgentsMcp { - fn drop(&mut self) { - self.task_manager.shutdown(); - } -} - struct InstrumentationGuard(Option>); impl InstrumentationGuard { diff --git a/crates/mcp-servers/src/subagents/tools/spawn_subagent/mod.rs b/crates/mcp-servers/src/subagents/tools/spawn_subagent/mod.rs index c873ba713..84786f095 100644 --- a/crates/mcp-servers/src/subagents/tools/spawn_subagent/mod.rs +++ b/crates/mcp-servers/src/subagents/tools/spawn_subagent/mod.rs @@ -1,10 +1,10 @@ -use crate::setup::McpBuilderExt; +use crate::setup::agent_mcp; use crate::subagents::SubAgentProgressSink; use aether_core::{ - agent_spec::McpConfigSource, - core::{AgentBuilder, AgentDeps, AgentHandle, recv_agent_event}, + agent_spec::AgentSpec, + core::{AgentBuilder, AgentDeps, AgentError, AgentHandle, AgentRegistryError, recv_agent_event}, events::{AgentEvent, Command, MessageEvent, SubAgentProgressPayload, TurnEvent, TurnOutcome, UserCommand}, - mcp::{McpRuntime, McpSession, mcp}, + mcp::{McpCatalog, McpSpawnError, ParseError}, }; use futures::FutureExt; use llm::{ContentBlock, MessageId}; @@ -12,8 +12,11 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::panic::AssertUnwindSafe; use std::path::PathBuf; +use std::time::Duration; +use thiserror::Error; use tokio::sync::mpsc; use tokio::task::JoinSet; +use tokio::time::timeout; use utils::display_meta::{ToolDisplayMeta, ToolResultMeta}; use utils::panic::panic_message; @@ -224,30 +227,24 @@ impl AgentExecutor { async fn execute_single(&self, task_id: String, task: SubAgentTask) -> SubAgentResult { let agent_name = task.agent_name.clone(); - let result: Result = async { - let spec = self - .deps - .agent_registry - .resolve_agent_invocable(&task.agent_name) - .map_err(|error| error.to_string())?; + let result: Result = async { + let spec = self.deps.agent_registry.resolve_agent_invocable(&task.agent_name)?; - let mut spawn = self.spawn_mcps(&spec.mcp_config_sources, spec.tools.clone()).await?; - spawn.block_until_ready().await.ok_or_else(|| "MCP bootstrap aborted before completion".to_string())?; - let mcp = spawn.handle().clone(); + let mcp = self.spawn_mcp(&spec)?; + if timeout(MCP_READY_TIMEOUT, mcp.ready()).await.is_err() { + tracing::warn!(agent = %agent_name, "Starting sub-agent before every MCP server connected"); + } - let (user_tx, mut agent_rx, agent_handle) = self.spawn_agent(spec, mcp).await?; - let (mcp_runtime, _) = spawn.connect_agent(user_tx.clone()).await.split(); - let running_agent = RunningAgent { user_tx, agent_handle, _mcp_runtime: mcp_runtime }; + let (user_tx, mut agent_rx, _agent) = self.spawn_agent(&spec, mcp).await?; let prompt_with_instructions = format!("{}\n\n{}", task.prompt, STRUCTURED_OUTPUT_INSTRUCTIONS); - running_agent - .user_tx + user_tx .send(Command::UserCommand(UserCommand::Text { message_id: MessageId::new(), content: vec![ContentBlock::text(&prompt_with_instructions)], })) .await - .map_err(|e| format!("Failed to send message to agent: {e}"))?; + .map_err(|_| SubAgentError::Stopped)?; self.report(&task_id, &agent_name, &AgentEvent::Turn(TurnEvent::Started { content: vec![] })).await; @@ -264,8 +261,8 @@ impl AgentExecutor { AgentEvent::Turn(TurnEvent::Ended { outcome }) => match outcome { TurnOutcome::Completed => return Ok(final_output), - TurnOutcome::Failed { error, .. } => return Err(format!("Agent error: {error}")), - TurnOutcome::Cancelled => return Err("Agent cancelled".to_string()), + TurnOutcome::Failed { error, .. } => return Err(SubAgentError::Failed(error.clone())), + TurnOutcome::Cancelled => return Err(SubAgentError::Cancelled), }, _ => {} @@ -282,9 +279,13 @@ impl AgentExecutor { output: Some(output), error: None, }, - Err(error) => { - SubAgentResult { task_id, agent_name, status: SubAgentStatus::Error, output: None, error: Some(error) } - } + Err(error) => SubAgentResult { + task_id, + agent_name, + status: SubAgentStatus::Error, + output: None, + error: Some(error.to_string()), + }, } } @@ -302,51 +303,43 @@ impl AgentExecutor { } } - async fn spawn_mcps( - &self, - effective_mcp_config_sources: &[McpConfigSource], - tool_filter: mcp_utils::client::ToolFilter, - ) -> Result { - let mut builder = mcp(&self.project_root) - .with_tool_filter(tool_filter) - .with_agent_deps(self.deps.clone()) - .with_builtin_servers(); - - if !effective_mcp_config_sources.is_empty() { - builder = builder - .from_mcp_config_sources(effective_mcp_config_sources) - .map_err(|e| format!("Failed to load mcp configs: {e}"))?; - } - - builder.spawn().await.map_err(|e| format!("Failed to spawn MCP manager: {e}")) + fn spawn_mcp(&self, spec: &AgentSpec) -> Result { + let builder = + agent_mcp(&self.project_root, self.deps.clone(), spec.tools.clone()).with_config(spec.mcp.clone())?; + Ok(builder.spawn()?) } async fn spawn_agent( &self, - spec: aether_core::agent_spec::AgentSpec, - mcp: aether_core::mcp::McpHandle, - ) -> Result<(mpsc::Sender, mpsc::Receiver, AgentHandle), String> { - AgentBuilder::from_spec(&spec, vec![], &self.deps) - .await - .map_err(|e| format!("Failed to build agent from spec: {e}"))? - .tools(mcp, Vec::new()) - .spawn() - .await - .map_err(|e| format!("Failed to spawn agent: {e}")) + spec: &AgentSpec, + mcp: McpCatalog, + ) -> Result<(mpsc::Sender, mpsc::Receiver, AgentHandle), SubAgentError> { + let builder = AgentBuilder::from_spec(spec, &self.deps).await.map_err(SubAgentError::Build)?; + builder.mcp(mcp).spawn().await.map_err(SubAgentError::Spawn) } } -struct RunningAgent { - user_tx: mpsc::Sender, - agent_handle: AgentHandle, - _mcp_runtime: McpRuntime, +#[derive(Debug, Error)] +enum SubAgentError { + #[error(transparent)] + Resolve(#[from] AgentRegistryError), + #[error("Failed to load MCP configs: {0}")] + McpConfig(#[from] ParseError), + #[error("Failed to spawn MCP servers: {0}")] + McpSpawn(#[from] McpSpawnError), + #[error("Failed to build agent from spec: {0}")] + Build(#[source] AgentError), + #[error("Failed to spawn agent: {0}")] + Spawn(#[source] AgentError), + #[error("Agent stopped before it received its task")] + Stopped, + #[error("Agent error: {0}")] + Failed(String), + #[error("Agent cancelled")] + Cancelled, } -impl Drop for RunningAgent { - fn drop(&mut self) { - self.agent_handle.abort(); - } -} +const MCP_READY_TIMEOUT: Duration = Duration::from_secs(30); #[cfg(test)] mod tests { diff --git a/crates/mcp-servers/src/tasks/server.rs b/crates/mcp-servers/src/tasks/server.rs index a6f02dda7..952226e36 100644 --- a/crates/mcp-servers/src/tasks/server.rs +++ b/crates/mcp-servers/src/tasks/server.rs @@ -1,3 +1,4 @@ +use crate::server_config; use clap::Parser; use rmcp::{ ServerHandler, @@ -5,14 +6,13 @@ use rmcp::{ router::tool::ToolRouter, wrapper::{Json, Parameters}, }, - model::{Implementation, ServerCapabilities, ServerConfig}, + model::{ServerCapabilities, ServerConfig}, tool, tool_handler, tool_router, }; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use tempfile::TempDir; -use crate::error::ServerInitError; use crate::{ tasks::{ TaskCreateInput, TaskCreateOutput, TaskGetInput, TaskGetOutput, TaskListInput, TaskListOutput, TaskStore, @@ -24,6 +24,7 @@ use crate::{ /// CLI arguments for `TasksMcp` server #[derive(Debug, Clone, Parser)] +#[command(name = "tasks", no_binary_name = true)] pub struct TasksMcpArgs { /// Base directory for persistent task storage. If omitted, tasks are /// session-scoped and stored in a temporary directory that is cleaned up @@ -32,15 +33,6 @@ pub struct TasksMcpArgs { pub dir: Option, } -impl TasksMcpArgs { - pub fn from_args(args: Vec) -> Result { - let mut full_args = vec!["tasks-mcp".to_string()]; - full_args.extend(args); - - Self::try_parse_from(full_args).map_err(ServerInitError::InvalidArgs) - } -} - #[doc = include_str!("../docs/tasks_mcp.md")] #[derive(Debug)] pub struct TasksMcp { @@ -84,24 +76,11 @@ impl TasksMcp { } } - /// Create a new `TasksMcp` server from parsed CLI arguments. - /// - /// If `--dir` is provided, tasks persist at that path. Otherwise, tasks are - /// session-scoped in a temporary directory. - pub fn from_args(args: Vec) -> Result { - let parsed_args = TasksMcpArgs::from_args(args)?; - Ok(match parsed_args.dir { - Some(dir) => Self::new_persistent(dir), - None => Self::new(), - }) - } - - pub fn from_args_with_base_dir(args: Vec, base_dir: &Path) -> Result { - let parsed_args = TasksMcpArgs::from_args(args)?; - Ok(match parsed_args.dir { + pub fn from_args(args: TasksMcpArgs, base_dir: &Path) -> Self { + match args.dir { Some(dir) => Self::new_persistent(resolve_path(base_dir, dir)), None => Self::new(), - }) + } } #[cfg(test)] @@ -121,9 +100,11 @@ impl TasksMcp { #[tool_handler(router = self.tool_router)] impl ServerHandler for TasksMcp { fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("tasks-mcp", "0.1.0")) - .with_instructions(include_str!("./instructions.md")) + server_config( + "tasks-mcp", + ServerCapabilities::builder().enable_tools().build(), + include_str!("./instructions.md"), + ) } } @@ -219,22 +200,17 @@ mod tests { } #[test] - fn test_from_args_no_dir_is_session_scoped() { - let server = TasksMcp::from_args(vec![]).unwrap(); - assert!(server.is_session_scoped()); - } - - #[test] - fn test_from_args_with_base_dir_resolves_relative_dir() { - let server = - TasksMcp::from_args_with_base_dir(vec!["--dir".into(), "tasks".into()], Path::new("/workspace")).unwrap(); + fn test_from_args_with_dir_is_persistent() { + let args = TasksMcpArgs::try_parse_from(["--dir", "tasks"]).unwrap(); + let server = TasksMcp::from_args(args, Path::new("/workspace")); assert!(!server.is_session_scoped()); } #[test] - fn test_from_args_with_base_dir_no_dir_stays_session_scoped() { - let server = TasksMcp::from_args_with_base_dir(vec![], Path::new("/workspace")).unwrap(); + fn test_from_args_no_dir_is_session_scoped() { + let server = + TasksMcp::from_args(TasksMcpArgs::try_parse_from(Vec::::new()).unwrap(), Path::new("/workspace")); assert!(server.is_session_scoped()); } } diff --git a/crates/mcp-servers/src/workspace_paths.rs b/crates/mcp-servers/src/workspace_paths.rs index 69656891b..1fcaa8397 100644 --- a/crates/mcp-servers/src/workspace_paths.rs +++ b/crates/mcp-servers/src/workspace_paths.rs @@ -44,6 +44,10 @@ pub fn resolve_path(root: &Path, path: PathBuf) -> PathBuf { if path.is_absolute() { path } else { root.join(path) } } +pub fn resolve_root(base: &Path, root: Option) -> PathBuf { + root.map_or_else(|| base.to_path_buf(), |path| resolve_path(base, path)) +} + pub fn relative_path(root: &Path, path: &str) -> String { Path::new(path) .strip_prefix(root) @@ -95,6 +99,13 @@ mod tests { ); } + #[test] + fn root_override_resolves_against_base() { + assert_eq!(resolve_root(Path::new("/workspace"), None), PathBuf::from("/workspace")); + assert_eq!(resolve_root(Path::new("/workspace"), Some("sub".into())), PathBuf::from("/workspace/sub")); + assert_eq!(resolve_root(Path::new("/workspace"), Some("/abs".into())), PathBuf::from("/abs")); + } + #[test] fn leaves_absolute_path_untouched() { assert_eq!(resolve_path(Path::new("/workspace"), PathBuf::from("/etc/passwd")), PathBuf::from("/etc/passwd")); diff --git a/crates/mcp-servers/tests/integration/builtin_servers.rs b/crates/mcp-servers/tests/integration/builtin_servers.rs new file mode 100644 index 000000000..a33ea8ac4 --- /dev/null +++ b/crates/mcp-servers/tests/integration/builtin_servers.rs @@ -0,0 +1,42 @@ +use aether_core::mcp::{McpConfig, McpSpawnError, mcp}; +use clap::ValueEnum; +use mcp_servers::{Builtin, McpBuilderExt}; +use tempfile::tempdir; + +#[tokio::test] +async fn invalid_permission_mode_fails_spawn_instead_of_disabling_permission_prompts() { + let error = spawn_error("coding", &["--permission-mode", "alwaysask"]); + + assert!(matches!(error, McpSpawnError::ServerFactory { ref name, .. } if name == "coding")); + assert!(error.to_string().contains("alwaysask"), "error should name the bad value: {error}"); +} + +#[tokio::test] +async fn every_builtin_server_rejects_unknown_args() { + for builtin in Builtin::value_variants() { + let error = spawn_error(builtin.name(), &["--unknown-flag"]); + + assert!(matches!(error, McpSpawnError::ServerFactory { ref name, .. } if name == builtin.name())); + } +} + +#[test] +fn builtin_servers_parse_from_their_names() { + for builtin in Builtin::value_variants() { + assert_eq!(Builtin::from_str(builtin.name(), false).unwrap(), *builtin); + } + + assert!(Builtin::from_str("nope", false).is_err()); +} + +fn spawn_error(server: &str, args: &[&str]) -> McpSpawnError { + let root = tempdir().expect("create workspace"); + let config = serde_json::json!({ "servers": { server: { "type": "in-memory", "args": args } } }); + let config = McpConfig::from_json(&config.to_string()).expect("config parses"); + let builder = mcp(root.path()).with_builtin_servers().with_config(config).expect("config is valid"); + + match builder.spawn() { + Ok(_) => panic!("{server} should reject {args:?}"), + Err(error) => error, + } +} diff --git a/crates/mcp-servers/tests/integration/coding_mcp_tools.rs b/crates/mcp-servers/tests/integration/coding_mcp_tools.rs index 1ac4e9b34..372300174 100644 --- a/crates/mcp-servers/tests/integration/coding_mcp_tools.rs +++ b/crates/mcp-servers/tests/integration/coding_mcp_tools.rs @@ -1,12 +1,13 @@ -use crate::common::{CodingWorkspace, TestResult, test_error}; -use aether_core::mcp::tool_bridge::convert_tool_result; -use llm::ToolCallRequest; +use crate::common::{CodingWorkspace, McpClientExt, TestResult, test_error}; +use aether_core::events::{AgentEvent, ToolEvent}; +use aether_core::testing::test_agent; +use llm::testing::llm_response; +use mcp_servers::CodingMcp; use mcp_servers::coding::tools::bash::BashInput; use mcp_servers::coding::tools::read_file::ReadFileArgs; use mcp_servers::coding::tools::web_fetch::{HttpResponse, WebFetchInput}; use mcp_servers::testing::FakeHttpClient; use std::fs::{canonicalize, read_to_string}; -use utils::temp_dir::TempDir; use utils::tool_result_truncator::saved_path; #[tokio::test] @@ -40,17 +41,28 @@ async fn full_read_file_pages_reach_the_llm_whole() -> TestResult { let workspace = CodingWorkspace::new().await?; let content = (1..=5_000).map(|n| format!("let line_{n} = \"{}\";", "x".repeat(20))).collect::>(); let path = workspace.write("big.rs", &content.join("\n"))?; - let page = workspace - .client - .call_raw("read_file", ReadFileArgs { file_path: path.to_string_lossy().into(), ..Default::default() }) - .await?; + let arguments = + serde_json::to_string(&ReadFileArgs { file_path: path.to_string_lossy().into(), ..Default::default() })?; - let request = ToolCallRequest { id: "read".into(), name: "coding__read_file".into(), arguments: "{}".into() }; - let (result, _) = - convert_tool_result(&request, Ok(page), &TempDir::new()).map_err(|error| test_error(error.error))?; + let events = test_agent() + .mcp_server("coding", CodingMcp::new(workspace.root())) + .llm_responses(&[ + llm_response().tool_call("read", "coding__read_file", &[&arguments]).build(), + llm_response().text(&["read it"]).build(), + ]) + .user_text("read the file") + .run() + .await?; - assert!(!result.result.contains("bytes omitted"), "{}", result.result); - assert!(result.result.contains("nextOffset"), "{}", result.result); + let result = events + .iter() + .find_map(|event| match event { + AgentEvent::Tool(ToolEvent::Result { result, .. }) => Some(result.result.as_str()), + _ => None, + }) + .ok_or_else(|| test_error("read_file did not return a result"))?; + assert!(!result.contains("bytes omitted"), "{result}"); + assert!(result.contains("nextOffset"), "{result}"); Ok(()) } @@ -174,8 +186,8 @@ async fn test_bash_pwd_uses_workspace_root() -> TestResult { #[tokio::test] async fn coding_tool_catalog_uses_bash_for_search_and_directory_listings() -> TestResult { let workspace = CodingWorkspace::new().await?; - let catalog = workspace.client.raw().list_tools(None).await?; - let names: Vec<_> = catalog.tools.iter().map(|tool| tool.name.as_ref()).collect(); + let catalog = workspace.client.list_tools().await?; + let names: Vec<_> = catalog.iter().map(|tool| tool.name.as_ref()).collect(); for retained in ["bash", "read_file", "write_file", "edit_file", "lsp_symbol"] { assert!(names.contains(&retained), "missing tool {retained}: {names:?}"); } diff --git a/crates/mcp-servers/tests/integration/coding_permissions.rs b/crates/mcp-servers/tests/integration/coding_permissions.rs index a1d743652..0adae34b6 100644 --- a/crates/mcp-servers/tests/integration/coding_permissions.rs +++ b/crates/mcp-servers/tests/integration/coding_permissions.rs @@ -1,12 +1,16 @@ -use crate::common::{TestClient, TestResult, scripted_mcp_client, silent_mcp_client, test_client_info}; +use crate::common::{McpClientExt, TestResult, cancels_prompts, connect, scripted, tool_request}; use mcp_servers::coding::CodingMcp; use mcp_servers::{DefaultCodingTools, PermissionMode}; -use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction}; +use mcp_utils::client::{ClientOptions, ElicitationSink}; +use mcp_utils::model::{ElicitationCapability, UrlElicitationCapability}; +use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction, ErrorCode, InputResponses}; +use rmcp::service::ServiceError; use serde_json::json; use tempfile::TempDir; +use tokio::sync::mpsc; fn coding_mcp(root: &TempDir, mode: PermissionMode) -> CodingMcp { - CodingMcp::new().with_root_dir(root.path().to_path_buf()).with_permission_mode(mode) + CodingMcp::new(root.path()).with_permission_mode(mode) } fn allow() -> ElicitResult { @@ -20,8 +24,8 @@ fn deny() -> ElicitResult { #[tokio::test] async fn always_ask_bash_runs_after_user_allows() -> TestResult { let root = TempDir::new()?; - let (client, elicitation) = scripted_mcp_client("coding", allow()); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), client).await?; + let (options, elicitation) = scripted(allow()); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; let result = mcp.call("bash", json!({ "command": "echo hello" })).await?; assert!(result["output"].as_str().unwrap_or_default().contains("hello")); @@ -38,8 +42,8 @@ async fn always_ask_bash_runs_after_user_allows() -> TestResult { #[tokio::test] async fn always_ask_bash_denied_returns_declined_error() -> TestResult { let root = TempDir::new()?; - let (client, _elicitation) = scripted_mcp_client("coding", deny()); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), client).await?; + let (options, _elicitation) = scripted(deny()); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; let result = mcp.call_raw("bash", json!({ "command": "echo hello" })).await?; assert!(result.is_error.unwrap_or(false), "denied command should error: {result:?}"); @@ -53,7 +57,7 @@ async fn auto_mode_safe_command_runs_without_elicitation() -> TestResult { let root = TempDir::new()?; // A silent client resolves any elicitation as Cancel, so success proves // no elicitation was dispatched. - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::Auto), silent_mcp_client("coding")).await?; + let mcp = connect(coding_mcp(&root, PermissionMode::Auto), cancels_prompts()).await; let result = mcp.call("bash", json!({ "command": "echo safe" })).await?; assert!(result["output"].as_str().unwrap_or_default().contains("safe")); @@ -63,8 +67,8 @@ async fn auto_mode_safe_command_runs_without_elicitation() -> TestResult { #[tokio::test] async fn auto_mode_dangerous_command_is_gated() -> TestResult { let root = TempDir::new()?; - let (client, elicitation) = scripted_mcp_client("coding", deny()); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::Auto), client).await?; + let (options, elicitation) = scripted(deny()); + let mcp = connect(coding_mcp(&root, PermissionMode::Auto), options).await; let result = mcp.call_raw("bash", json!({ "command": "rm -rf ./scratch" })).await?; assert!(result.is_error.unwrap_or(false), "denied dangerous command should error: {result:?}"); @@ -80,8 +84,8 @@ async fn auto_mode_dangerous_command_is_gated() -> TestResult { #[tokio::test] async fn always_ask_write_file_denied_leaves_file_unwritten() -> TestResult { let root = TempDir::new()?; - let (client, _elicitation) = scripted_mcp_client("coding", deny()); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), client).await?; + let (options, _elicitation) = scripted(deny()); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; let target = root.path().join("never.txt"); let result = @@ -94,7 +98,7 @@ async fn always_ask_write_file_denied_leaves_file_unwritten() -> TestResult { #[tokio::test] async fn auto_mode_write_file_is_not_gated() -> TestResult { let root = TempDir::new()?; - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::Auto), silent_mcp_client("coding")).await?; + let mcp = connect(coding_mcp(&root, PermissionMode::Auto), cancels_prompts()).await; let target = root.path().join("free.txt"); mcp.call("write_file", json!({ "filePath": target.display().to_string(), "content": "hi" })).await?; @@ -106,8 +110,8 @@ async fn auto_mode_write_file_is_not_gated() -> TestResult { async fn cancelled_prompt_carrying_allow_content_is_not_approval() -> TestResult { let root = TempDir::new()?; let cancel_with_allow = ElicitResult::new(ElicitationAction::Cancel).with_content(json!({ "decision": "allow" })); - let (client, _elicitation) = scripted_mcp_client("coding", cancel_with_allow); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), client).await?; + let (options, _elicitation) = scripted(cancel_with_allow); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; let target = root.path().join("never.txt"); let result = mcp.call_raw("bash", json!({ "command": format!("touch {}", target.display()) })).await?; @@ -119,8 +123,8 @@ async fn cancelled_prompt_carrying_allow_content_is_not_approval() -> TestResult #[tokio::test] async fn always_ask_gates_lsp_rename_via_destructive_annotation() -> TestResult { let root = TempDir::new()?; - let (client, elicitation) = scripted_mcp_client("coding", deny()); - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), client).await?; + let (options, elicitation) = scripted(deny()); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; let result = mcp.call_raw("lsp_rename", json!({ "symbol": "foo", "newName": "bar" })).await?; assert!(result.is_error.unwrap_or(false), "denied rename should error: {result:?}"); @@ -138,7 +142,7 @@ async fn always_ask_gates_lsp_rename_via_destructive_annotation() -> TestResult #[tokio::test] async fn always_ask_without_elicitation_capability_denies_with_friendly_error() -> TestResult { let root = TempDir::new()?; - let mcp = TestClient::start_with(|| coding_mcp(&root, PermissionMode::AlwaysAsk), test_client_info()).await?; + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), ClientOptions::default()).await; let target = root.path().join("gated.txt"); let result = mcp.call_raw("bash", json!({ "command": format!("touch {}", target.display()) })).await?; @@ -148,3 +152,50 @@ async fn always_ask_without_elicitation_capability_denies_with_friendly_error() assert!(!target.exists(), "the gated command must not run"); Ok(()) } + +#[tokio::test] +async fn clients_that_only_open_urls_cannot_answer_permission_prompts() -> TestResult { + let root = TempDir::new()?; + let (events, _host) = mpsc::channel(1); + let url_only = ElicitationCapability::new().with_url(UrlElicitationCapability::new()); + let options = ClientOptions::default().elicitation(ElicitationSink::new(events, url_only)); + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), options).await; + + let result = mcp.call_raw("bash", json!({ "command": "echo hello" })).await?; + + let text = result.content.first().and_then(|c| c.as_text()).expect("error text"); + assert!(text.text.contains("cannot prompt"), "url-only clients cannot fill in forms: {}", text.text); + Ok(()) +} + +#[tokio::test] +async fn permission_answers_must_include_the_permission_decision() -> TestResult { + let root = TempDir::new()?; + let mcp = connect(coding_mcp(&root, PermissionMode::AlwaysAsk), cancels_prompts()).await; + + let missing = answer_permission(&mcp, InputResponses::new()).await; + let malformed = answer_permission(&mcp, InputResponses::from([("permission".to_string(), json!("yes"))])).await; + + assert_invalid_params(missing, "missing input response for 'permission'"); + assert_invalid_params(malformed, "invalid input response for 'permission'"); + Ok(()) +} + +async fn answer_permission( + mcp: &mcp_utils::client::McpClient, + responses: InputResponses, +) -> Result { + let mut request = tool_request("bash", json!({ "command": "echo hello" })); + request.input_responses = Some(responses); + mcp.peer().call_tool_once(request).await +} + +fn assert_invalid_params(result: Result, message: &str) { + match result { + Err(ServiceError::McpError(error)) => { + assert_eq!(error.code, ErrorCode::INVALID_PARAMS); + assert!(error.message.contains(message), "unexpected error: {}", error.message); + } + other => panic!("expected an invalid params error, got {other:?}"), + } +} diff --git a/crates/mcp-servers/tests/integration/common/mod.rs b/crates/mcp-servers/tests/integration/common/mod.rs index 5231ea0a8..8e449587e 100644 --- a/crates/mcp-servers/tests/integration/common/mod.rs +++ b/crates/mcp-servers/tests/integration/common/mod.rs @@ -10,50 +10,74 @@ use mcp_servers::skills::{ tools::{LoadSkillsInput, SkillRequest}, }; use mcp_servers::testing::FakeHttpClient; -use mcp_utils::client::{McpClient, client_capabilities}; -use mcp_utils::testing::{ElicitationScript, connect}; -use rmcp::RoleClient; +use mcp_utils::client::{ClientOptions, McpClient}; +use mcp_utils::testing::{ElicitationScript, args, call, elicitation_host}; +use rmcp::ServerHandler; use rmcp::model::{ - CallToolRequestParams, CallToolResult, ClientCapabilities, ClientConfig, ElicitResult, Implementation, + CallToolRequestParams, CallToolResponse, CallToolResult, DetailedTask, ElicitResult, GetTaskParams, TaskPayload, + TaskStatus, }; -use rmcp::service::RunningService; -use rmcp::{RoleServer, Service}; use serde::Serialize; +use serde_json::Value; use std::fs::{create_dir_all, read_to_string}; use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; use tempfile::{TempDir, tempdir}; -use tokio::sync::mpsc; /// Default timeout for polling operations (60 seconds). const POLL_TIMEOUT: Duration = Duration::from_mins(1); const POLL_INTERVAL: Duration = Duration::from_millis(500); +const TASK_POLL_INTERVAL: Duration = Duration::from_millis(10); pub type TestResult = Result>; -pub fn test_client_info() -> ClientConfig { - ClientConfig::new(ClientCapabilities::default(), Implementation::new("test-client", "0.1.0")) +pub async fn connect(server: impl ServerHandler, options: ClientOptions) -> McpClient { + mcp_utils::testing::connect("test", server, &options).await } -/// Client info declaring the same capabilities `McpManager` advertises in -/// production: form + url elicitation and MCP tasks. -pub fn production_client_info() -> ClientConfig { - ClientConfig::new(client_capabilities(), Implementation::new("test-client", "0.1.0")) +pub fn cancels_prompts() -> ClientOptions { + elicitation_host().0 } -/// An `McpClient` whose elicitation events go nowhere; elicitation-dependent -/// calls resolve as Cancel. -pub fn silent_mcp_client(server_name: &str) -> McpClient { - let (event_tx, _event_rx) = mpsc::channel(8); - McpClient::new(production_client_info(), server_name.to_string(), event_tx) +/// A client that answers its first elicitation with `response`. +pub fn scripted(response: ElicitResult) -> (ClientOptions, ElicitationScript) { + ElicitationScript::answering([response]) } -/// An `McpClient` plus a script answering the next elicitation request with -/// `response` and capturing what arrived for assertions. -pub fn scripted_mcp_client(server_name: &str, response: ElicitResult) -> (McpClient, ElicitationScript) { - let (event_tx, event_rx) = mpsc::channel(8); - let client = McpClient::new(production_client_info(), server_name.to_string(), event_tx); - (client, ElicitationScript::spawn(event_rx, [response])) +pub fn tool_request(tool: &str, arguments: impl Serialize) -> CallToolRequestParams { + let arguments = serde_json::to_value(arguments).expect("tool arguments serialize to JSON"); + CallToolRequestParams::new(tool.to_string()).with_arguments(args(arguments)) +} + +pub fn taskless() -> ClientOptions { + ClientOptions::default().without_tasks() +} + +pub fn coding_with_lsp(project: &impl TestProject) -> CodingMcp { + CodingMcp::new(project.root()).with_lsp() +} + +pub trait McpClientExt { + async fn call_raw(&self, tool: &str, args: impl Serialize) -> TestResult; + + async fn call(&self, tool: &str, args: impl Serialize) -> TestResult; +} + +impl McpClientExt for McpClient { + async fn call_raw(&self, tool: &str, args: impl Serialize) -> TestResult { + Ok(call(self, tool, serde_json::to_value(args)?).await?) + } + + async fn call(&self, tool: &str, args: impl Serialize) -> TestResult { + let result = self.call_raw(tool, args).await?; + let text = result + .content + .first() + .and_then(|content| content.as_text()) + .ok_or_else(|| test_error(format!("{tool} returned no text content: {result:?}")))?; + serde_json::from_str(&text.text) + .map_err(|error| test_error(format!("{tool} returned non-JSON text ({error}): {}", text.text)).into()) + } } pub fn test_error(message: impl Into) -> std::io::Error { @@ -63,28 +87,25 @@ pub fn test_error(message: impl Into) -> std::io::Error { /// An isolated workspace connected to a `CodingMcp` through the public MCP protocol. pub struct CodingWorkspace { root: TempDir, - pub client: TestClient, + pub client: McpClient, } impl CodingWorkspace { pub async fn new() -> TestResult { - Self::start(|root| CodingMcp::new().with_root_dir(root.to_path_buf())).await + Self::start(|root| CodingMcp::new(root)).await } pub async fn new_with_lsp() -> TestResult { - Self::start(|root| CodingMcp::new().with_lsp(root.to_path_buf())).await + Self::start(|root| CodingMcp::new(root).with_lsp()).await } pub async fn with_http(http: FakeHttpClient) -> TestResult { - Self::start(|root| { - CodingMcp::new().with_root_dir(root.to_path_buf()).with_web_fetcher(WebFetcher::with_client(http)) - }) - .await + Self::start(|root| CodingMcp::new(root).with_web_fetcher(WebFetcher::with_client(http))).await } async fn start(configure: impl FnOnce(&Path) -> CodingMcp) -> TestResult { let root = tempdir()?; - let client = TestClient::start(|| configure(root.path())).await?; + let client = connect(configure(root.path()), ClientOptions::default()).await; Ok(Self { root, client }) } @@ -110,64 +131,8 @@ impl CodingWorkspace { } } -/// Generic test client wrapping a connected MCP server. -pub struct TestClient, U: Service = ClientConfig> { - _server_handle: RunningService, - client: RunningService, -} - -impl> TestClient { - /// Connect a default test `ClientConfig` to a server built by `configure`. - pub async fn start(configure: impl FnOnce() -> T) -> TestResult { - let (server_handle, client) = connect(configure(), test_client_info()).await?; - Ok(Self { _server_handle: server_handle, client }) - } -} - -impl, U: Service> TestClient { - pub async fn start_with(configure_server: impl FnOnce() -> T, client: U) -> TestResult { - let (server_handle, client) = connect(configure_server(), client).await?; - Ok(Self { _server_handle: server_handle, client }) - } - - pub async fn call(&self, tool: &str, args: V) -> TestResult { - let result = self.call_raw(tool, args).await?; - let text = result - .content - .first() - .and_then(|c| c.as_text()) - .ok_or_else(|| test_error(format!("{tool} should return text content")))?; - Ok(serde_json::from_str(&text.text)?) - } - - pub async fn call_raw(&self, tool: &str, args: V) -> TestResult { - let args = serde_json::to_value(args)?; - let arguments = - args.as_object().ok_or_else(|| test_error("tool arguments must serialize to a JSON object"))?.clone(); - Ok(self.client.call_tool(CallToolRequestParams::new(tool.to_string()).with_arguments(arguments)).await?) - } - - pub fn raw(&self) -> &RunningService { - &self.client - } -} - -pub async fn connect_lsp( - project: &impl TestProject, -) -> (RunningService, RunningService) { - let server = CodingMcp::new().with_lsp(project.root().to_path_buf()); - connect(server, test_client_info()).await.expect("Failed to connect") -} - -pub async fn call_tool_error( - client: &RunningService, - name: &str, - args: serde_json::Value, -) -> String { - match client - .call_tool(CallToolRequestParams::new(name.to_string()).with_arguments(args.as_object().unwrap().clone())) - .await - { +pub async fn call_tool_error(client: &McpClient, name: &str, args: Value) -> String { + match client.call_raw(name, args).await { Ok(result) => { assert!(result.is_error.unwrap_or(false), "tool call should fail: {result:?}"); let content = result.content.first().expect("Expected error content"); @@ -178,43 +143,40 @@ pub async fn call_tool_error( } } -pub async fn call_tool( - client: &RunningService, - name: &str, - args: serde_json::Value, -) -> serde_json::Value { - try_call_tool(client, name, args).await.unwrap_or_else(|| panic!("Tool '{name}' did not return valid JSON")) +pub async fn call_tool(client: &McpClient, name: &str, args: Value) -> Value { + client.call(name, args).await.unwrap_or_else(|error| panic!("Tool '{name}' did not return valid JSON: {error}")) } -pub async fn try_call_tool( - client: &RunningService, - name: &str, - args: serde_json::Value, -) -> Option { - let result = match client - .call_tool(CallToolRequestParams::new(name.to_string()).with_arguments(args.as_object().unwrap().clone())) - .await - { - Ok(result) => result, - Err(error) => { - eprintln!("[try_call_tool] {name} RPC error: {error}"); - return None; - } +pub async fn try_call_tool(client: &McpClient, name: &str, args: Value) -> Option { + client.call(name, args).await.inspect_err(|error| eprintln!("[try_call_tool] {name}: {error}")).ok() +} + +pub async fn call_task(client: &McpClient, tool: &str, arguments: Value) -> TestResult { + let response = client.peer().call_tool_once(tool_request(tool, arguments)).await?; + let CallToolResponse::Task(created) = response else { + return Err(test_error(format!("expected task response: {response:?}")).into()); }; - let Some(text) = result.content.first().and_then(|c| c.as_text()) else { - eprintln!("[try_call_tool] {name} no text content in response"); - return None; + assert_eq!(created.task.status, TaskStatus::Working); + + let task = await_terminal_task(client, &created.task.task_id).await?; + let TaskPayload::Completed { result } = task.payload else { + return Err(test_error(format!("expected completed task: {task:?}")).into()); }; - if let Ok(value) = serde_json::from_str(&text.text) { - Some(value) - } else { - eprintln!("[try_call_tool] {name} non-JSON response: {}", text.text); - None + Ok(serde_json::from_value(Value::Object(result))?) +} + +pub async fn await_terminal_task(client: &McpClient, task_id: &str) -> TestResult { + loop { + let task = client.peer().get_task(GetTaskParams::new(task_id.to_string())).await?.task; + if task.status().is_terminal() { + return Ok(task); + } + tokio::time::sleep(TASK_POLL_INTERVAL).await; } } pub async fn poll_diagnostics( - client: &RunningService, + client: &McpClient, file_path: Option<&str>, predicate: impl Fn(&serde_json::Value) -> bool, ) -> serde_json::Value { @@ -226,7 +188,7 @@ pub async fn poll_diagnostics( } pub async fn poll_lsp_tool( - client: &RunningService, + client: &McpClient, tool_name: &str, args: serde_json::Value, predicate: impl Fn(&serde_json::Value) -> bool, @@ -286,7 +248,7 @@ pub fn create_test_files(files: &[(&str, &str)]) -> TempDir { /// A skills server serving the `skills` directory of `test_dir`. pub fn skills_server(test_dir: &Path) -> SkillsMcp { - SkillsMcp::new(&[test_dir.join("skills")]) + SkillsMcp::new(test_dir, &[test_dir.join("skills")]) } pub fn load_skills_input(requests: &[(&str, Option<&str>)]) -> LoadSkillsInput { diff --git a/crates/mcp-servers/tests/integration/lsp_check_errors_contract.rs b/crates/mcp-servers/tests/integration/lsp_check_errors_contract.rs index 49bc9af54..cfe87efc1 100644 --- a/crates/mcp-servers/tests/integration/lsp_check_errors_contract.rs +++ b/crates/mcp-servers/tests/integration/lsp_check_errors_contract.rs @@ -1,5 +1,6 @@ -use crate::common::{CodingWorkspace, call_tool, call_tool_error, connect_lsp}; +use crate::common::{CodingWorkspace, call_tool, call_tool_error, coding_with_lsp, connect}; use aether_lspd::testing::{CargoProject, TestProject}; +use mcp_utils::client::ClientOptions; use std::os::unix::fs::PermissionsExt; #[tokio::test] @@ -7,7 +8,7 @@ async fn lsp_check_errors_accepts_flat_file_path_and_infers_file_scope() { let project = CargoProject::new("diag_contract_infers_file_scope").expect("Failed to create project"); project.add_file("src/main.rs", "fn main() {}\n").expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = call_tool( &client, "lsp_check_errors", @@ -25,8 +26,8 @@ async fn lsp_check_errors_schema_is_flat_and_has_only_file_path() { let project = CargoProject::new("diag_contract_flat_schema").expect("Failed to create project"); project.add_file("src/main.rs", "fn main() {}\n").expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; - let tools = client.peer().list_all_tools().await.expect("list tools"); + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; + let tools = client.list_tools().await.expect("list tools"); let tool = tools.into_iter().find(|tool| tool.name.as_ref() == "lsp_check_errors").expect("lsp_check_errors tool present"); @@ -44,7 +45,7 @@ async fn lsp_check_errors_rejects_redundant_scope_parameter() { let project = CargoProject::new("diag_contract_redundant_scope").expect("Failed to create project"); project.add_file("src/main.rs", "fn main() {}\n").expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let error = call_tool_error(&client, "lsp_check_errors", serde_json::json!({ "scope": "workspace" })).await; assert!(error.contains("unknown field `scope`"), "{error}"); @@ -53,7 +54,7 @@ async fn lsp_check_errors_rejects_redundant_scope_parameter() { #[tokio::test] async fn lsp_check_errors_fails_when_workspace_has_no_active_language_server() { let workspace = CodingWorkspace::new_with_lsp().await.expect("create workspace"); - let error = call_tool_error(workspace.client.raw(), "lsp_check_errors", serde_json::json!({})).await; + let error = call_tool_error(&workspace.client, "lsp_check_errors", serde_json::json!({})).await; assert!(error.contains("No active LSP clients"), "{error}"); } @@ -66,7 +67,7 @@ async fn lsp_check_errors_returns_typescript_installation_instructions_when_serv workspace.write("node_modules/.bin/tsc", "not executable\n").expect("write unavailable language server"); let error = - call_tool_error(workspace.client.raw(), "lsp_check_errors", serde_json::json!({ "filePath": index_ts })).await; + call_tool_error(&workspace.client, "lsp_check_errors", serde_json::json!({ "filePath": index_ts })).await; assert!(error.contains("Permission denied"), "{error}"); assert!(error.contains("TypeScript 7 language server"), "{error}"); @@ -85,7 +86,7 @@ async fn lsp_check_errors_returns_typescript_installation_instructions_when_tsc_ std::fs::set_permissions(&tsc, std::fs::Permissions::from_mode(0o755)).expect("make tsc executable"); let error = - call_tool_error(workspace.client.raw(), "lsp_check_errors", serde_json::json!({ "filePath": index_ts })).await; + call_tool_error(&workspace.client, "lsp_check_errors", serde_json::json!({ "filePath": index_ts })).await; assert!(error.contains("failed to complete initialization"), "{error}"); assert!(error.contains("TypeScript 7 language server"), "{error}"); @@ -97,7 +98,7 @@ async fn lsp_check_errors_rejects_file_scope_parameter() { let project = CargoProject::new("diag_contract_rejects_scope").expect("Failed to create project"); project.add_file("src/main.rs", "fn main() {}\n").expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let error = call_tool_error( &client, "lsp_check_errors", @@ -116,7 +117,7 @@ async fn lsp_check_errors_rejects_file_scope_directory_path() { let project = CargoProject::new("diag_contract_file_rejects_directory").expect("Failed to create project"); project.add_file("src/main.rs", "fn main() {}\n").expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let error = call_tool_error( &client, "lsp_check_errors", diff --git a/crates/mcp-servers/tests/integration/lsp_diagnostics_e2e.rs b/crates/mcp-servers/tests/integration/lsp_diagnostics_e2e.rs index 961bafc68..0eb04ec33 100644 --- a/crates/mcp-servers/tests/integration/lsp_diagnostics_e2e.rs +++ b/crates/mcp-servers/tests/integration/lsp_diagnostics_e2e.rs @@ -9,16 +9,14 @@ //! //! Run with: `cargo test -p mcp-servers -- lsp_diagnostics` -use crate::common::{call_tool, connect_lsp, has_errors, has_no_errors, poll_diagnostics, try_call_tool}; +use crate::common::{call_tool, coding_with_lsp, connect, has_errors, has_no_errors, poll_diagnostics, try_call_tool}; use aether_lspd::testing::{CargoProject, TestProject}; -use rmcp::RoleClient; -use rmcp::model::ClientConfig; -use rmcp::service::RunningService; +use mcp_utils::client::{ClientOptions, McpClient}; use std::path::PathBuf; use std::time::{Duration, Instant}; async fn poll_workspace_diagnostics( - client: &RunningService, + client: &McpClient, predicate: impl Fn(&serde_json::Value) -> bool, timeout: Duration, ) -> serde_json::Value { @@ -76,7 +74,7 @@ async fn test_mcp_edit_produces_diagnostics() { let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_diagnostics(&client, Some(&main_rs), has_errors).await; let errors = result["summary"]["errors"].as_u64().unwrap(); @@ -132,7 +130,7 @@ async fn test_diagnostics_available_after_edit_without_polling() { let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, Some(&main_rs), has_no_errors).await; @@ -170,7 +168,7 @@ async fn test_diagnostics_all_files_after_edit() { let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, Some(&main_rs), has_no_errors).await; @@ -208,7 +206,7 @@ async fn test_workspace_diagnostics_after_edit_without_file_check() { let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, None, has_no_errors).await; @@ -246,7 +244,7 @@ async fn test_workspace_diagnostics_clear_after_fix_without_file_check() { let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let initial = poll_workspace_diagnostics(&client, |result| file_error_count(result, &main_rs) > 0, Duration::from_mins(1)) @@ -299,7 +297,7 @@ async fn test_diagnostics_all_files_after_external_edit() { let main_rs = project.file_path_str("src/main.rs"); let main_rs_path = project.root().join("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, Some(&main_rs), has_no_errors).await; @@ -336,7 +334,7 @@ async fn test_diagnostics_all_files_after_external_edit_single_call() { let main_rs = project.file_path_str("src/main.rs"); let main_rs_path = project.root().join("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, Some(&main_rs), has_no_errors).await; @@ -371,7 +369,7 @@ async fn test_external_file_change_produces_diagnostics() { let main_rs = project.file_path_str("src/main.rs"); let main_rs_path = project.root().join("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_diagnostics(&client, Some(&main_rs), has_errors).await; let errors = result["summary"]["errors"].as_u64().unwrap(); @@ -429,7 +427,7 @@ async fn test_diagnostics_all_files_discovers_file_watcher_uris() { let main_rs_path = project.root().join("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; poll_diagnostics(&client, None, has_no_errors).await; diff --git a/crates/mcp-servers/tests/integration/lsp_operations_e2e.rs b/crates/mcp-servers/tests/integration/lsp_operations_e2e.rs index 2ba5cad31..1fa721c6b 100644 --- a/crates/mcp-servers/tests/integration/lsp_operations_e2e.rs +++ b/crates/mcp-servers/tests/integration/lsp_operations_e2e.rs @@ -7,8 +7,9 @@ //! //! Run with: `cargo test -p mcp-servers -- lsp_operations` -use crate::common::{call_tool, connect_lsp, poll_lsp_tool}; +use crate::common::{call_tool, coding_with_lsp, connect, poll_lsp_tool}; use aether_lspd::testing::{CargoProject, TestProject}; +use mcp_utils::client::ClientOptions; /// Test: hover returns type information for a Rust variable #[tokio::test] @@ -26,7 +27,7 @@ async fn test_hover_falls_back_from_stale_line_hint() { .expect("Failed to add file"); let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, @@ -65,7 +66,7 @@ fn main() { .expect("Failed to add file"); let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, @@ -104,7 +105,7 @@ fn main() { .expect("Failed to add file"); let main_rs = project.file_path_str("src/main.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, @@ -152,7 +153,7 @@ async fn test_lsp_rename_falls_back_from_stale_line_hint() { .expect("Failed to add main.rs"); let lib_rs = project.file_path_str("src/lib.rs"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; // Wait for rust-analyzer to fully index both files before calling rename. // Rename is a side-effecting operation that writes to disk, so it must not diff --git a/crates/mcp-servers/tests/integration/lsp_ts_diagnostics_e2e.rs b/crates/mcp-servers/tests/integration/lsp_ts_diagnostics_e2e.rs index 02c5d064d..7bac76f45 100644 --- a/crates/mcp-servers/tests/integration/lsp_ts_diagnostics_e2e.rs +++ b/crates/mcp-servers/tests/integration/lsp_ts_diagnostics_e2e.rs @@ -1,8 +1,6 @@ -use crate::common::{call_tool, connect_lsp, has_errors, has_no_errors}; +use crate::common::{call_tool, coding_with_lsp, connect, has_errors, has_no_errors}; use aether_lspd::testing::{NodeProject, TestProject}; -use rmcp::RoleClient; -use rmcp::model::ClientConfig; -use rmcp::service::RunningService; +use mcp_utils::client::{ClientOptions, McpClient}; #[tokio::test] async fn test_ts_mcp_edit_produces_diagnostics() { @@ -13,7 +11,7 @@ async fn test_ts_mcp_edit_produces_diagnostics() { let index_ts = project.file_path_str("src/index.ts"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = check_errors(&client, &index_ts).await; assert!(has_errors(&result), "Expected type error diagnostics: {result}"); @@ -59,7 +57,7 @@ async fn test_ts_external_file_change_produces_diagnostics() { let index_ts = project.file_path_str("src/index.ts"); let index_ts_path = project.root().join("src/index.ts"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = check_errors(&client, &index_ts).await; assert!(has_errors(&result), "Expected type error diagnostics: {result}"); @@ -82,7 +80,7 @@ async fn test_ts_diagnostics_after_edit_without_polling() { let index_ts = project.file_path_str("src/index.ts"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; call_tool(&client, "read_file", serde_json::json!({ "filePath": index_ts })).await; @@ -100,6 +98,6 @@ async fn test_ts_diagnostics_after_edit_without_polling() { assert!(has_errors(&result), "Expected diagnostics after edit + single lsp_check_errors call: {result}"); } -async fn check_errors(client: &RunningService, file_path: &str) -> serde_json::Value { +async fn check_errors(client: &McpClient, file_path: &str) -> serde_json::Value { call_tool(client, "lsp_check_errors", serde_json::json!({ "filePath": file_path })).await } diff --git a/crates/mcp-servers/tests/integration/lsp_ts_operations_e2e.rs b/crates/mcp-servers/tests/integration/lsp_ts_operations_e2e.rs index 8f553ab30..aa509db1c 100644 --- a/crates/mcp-servers/tests/integration/lsp_ts_operations_e2e.rs +++ b/crates/mcp-servers/tests/integration/lsp_ts_operations_e2e.rs @@ -1,5 +1,6 @@ -use crate::common::{connect_lsp, poll_lsp_tool}; +use crate::common::{coding_with_lsp, connect, poll_lsp_tool}; use aether_lspd::testing::{NodeProject, TestProject}; +use mcp_utils::client::ClientOptions; #[tokio::test] async fn test_ts_hover_returns_type_info() { @@ -7,7 +8,7 @@ async fn test_ts_hover_returns_type_info() { project.add_file("src/index.ts", "const x: number = 42;\nconsole.log(x);\n").expect("Failed to add file"); let index_ts = project.file_path_str("src/index.ts"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, @@ -42,7 +43,7 @@ console.log(msg); ) .expect("Failed to add file"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, @@ -78,7 +79,7 @@ console.log(a, b); .expect("Failed to add file"); let index_ts = project.file_path_str("src/index.ts"); - let (_server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let result = poll_lsp_tool( &client, diff --git a/crates/mcp-servers/tests/integration/main.rs b/crates/mcp-servers/tests/integration/main.rs index f57ee7302..4f368b59c 100644 --- a/crates/mcp-servers/tests/integration/main.rs +++ b/crates/mcp-servers/tests/integration/main.rs @@ -1,5 +1,6 @@ mod common; +mod builtin_servers; mod coding_mcp_tools; mod coding_permissions; mod lsp_check_errors_contract; diff --git a/crates/mcp-servers/tests/integration/plugins_mcp_agents.rs b/crates/mcp-servers/tests/integration/plugins_mcp_agents.rs index 233d6dc0d..21767c470 100644 --- a/crates/mcp-servers/tests/integration/plugins_mcp_agents.rs +++ b/crates/mcp-servers/tests/integration/plugins_mcp_agents.rs @@ -1,23 +1,22 @@ -use crate::common::{TestClient, TestResult, create_test_files, production_client_info, test_client_info, test_error}; +use crate::common::{ + TestResult, await_terminal_task, call_task, cancels_prompts, connect, create_test_files, taskless, test_error, + tool_request, +}; use aether_auth::FakeOAuthCredentialStore; -use aether_core::agent_spec::McpConfigSource; use aether_core::core::AgentDeps; use aether_core::events::{AgentEvent, AgentObserver, McpRequestInstrumentation, ObserverFactory, TraceContext}; -use aether_core::mcp::mcp; -use aether_core::mcp::tool_bridge::convert_tool_result; +use aether_core::mcp::{McpBuilder, McpConfig, ToolFilter}; +use aether_core::testing::ready_catalog; use aether_project::{AetherSettings, AetherSettingsSource, AgentCatalog, SettingsFileSource}; use futures::StreamExt; -use mcp_servers::McpBuilderExt; +use mcp_servers::agent_mcp; use mcp_servers::subagents::SubAgentsMcp; use mcp_servers::subagents::tools::{SpawnSubAgentsInput, SubAgentTask}; -use mcp_utils::client::{CallToolOptions, CancellationToken, ToolCallEvent}; -use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, DetailedTask, GetTaskParams, - TaskPayload, TaskStatus, -}; +use mcp_utils::client::{ClientOptions, McpClient, ToolCallEvent, ToolCallOptions}; +use mcp_utils::testing::args; +use rmcp::model::{CallToolResponse, CallToolResult, CancelTaskParams, GetTaskParams, TaskPayload}; use std::path::Path; use std::sync::{Arc, Mutex}; -use std::time::Duration; use tempfile::TempDir; fn spawn_input(tasks: &[(&str, &str)]) -> SpawnSubAgentsInput { @@ -136,21 +135,15 @@ async fn test_spawn_agent_with_coding_mcp_from_settings_catalog() { ]; let temp_dir = create_test_files(&test_files); - let _mcp = TestClient::start(|| create_test_server(temp_dir.path())).await.unwrap(); + let _mcp = connect(create_test_server(temp_dir.path()), ClientOptions::default()).await; } #[tokio::test] async fn test_spawn_subagent_codex_uses_oauth_store() -> TestResult { let temp_dir = create_project_with_codex_agent(); - let mcp = TestClient::start_with( - || { - let deps = AgentDeps::new(Arc::new(FakeOAuthCredentialStore::new()), None) - .with_agent_registry(test_registry(temp_dir.path())); - SubAgentsMcp::embedded(temp_dir.path().to_path_buf(), deps) - }, - production_client_info(), - ) - .await?; + let deps = AgentDeps::new(Arc::new(FakeOAuthCredentialStore::new()), None) + .with_agent_registry(test_registry(temp_dir.path())); + let mcp = connect(SubAgentsMcp::new(temp_dir.path().to_path_buf(), deps), cancels_prompts()).await; let parsed = call_complete(&mcp, spawn_input(&[("explorer", "Read README.md")])).await?; @@ -162,7 +155,7 @@ async fn test_spawn_subagent_codex_uses_oauth_store() -> TestResult { #[tokio::test] async fn spawn_subagent_empty_batch_completes_with_empty_output() -> TestResult { let temp_dir = TempDir::new().expect("Failed to create temp directory"); - let mcp = TestClient::start_with(|| create_test_server(temp_dir.path()), test_client_info()).await?; + let mcp = connect(create_test_server(temp_dir.path()), taskless()).await; let parsed = call_complete(&mcp, spawn_input(&[])).await?; @@ -176,10 +169,13 @@ async fn spawn_subagent_empty_batch_completes_with_empty_output() -> TestResult #[tokio::test] async fn test_spawn_subagent_errors_when_no_invocable_agents() -> TestResult { let temp_dir = TempDir::new().expect("Failed to create temp directory"); - let mcp = TestClient::start_with(|| create_test_server(temp_dir.path()), production_client_info()).await?; + let mcp = connect(create_test_server(temp_dir.path()), cancels_prompts()).await; - let error = - mcp.raw().call_tool_once(tool_request(spawn_input(&[("any-agent", "Do something")]))).await.unwrap_err(); + let error = mcp + .peer() + .call_tool_once(tool_request("spawn_subagent", spawn_input(&[("any-agent", "Do something")]))) + .await + .unwrap_err(); assert!( error.to_string().contains("No agent-invocable sub-agents are registered"), @@ -191,7 +187,7 @@ async fn test_spawn_subagent_errors_when_no_invocable_agents() -> TestResult { #[tokio::test] async fn spawn_subagent_defaults_to_foreground_without_tasks_capability() -> TestResult { let temp_dir = create_project_with_invocable_agent(); - let mcp = TestClient::start_with(|| create_test_server(temp_dir.path()), test_client_info()).await?; + let mcp = connect(create_test_server(temp_dir.path()), taskless()).await; let parsed = call_complete(&mcp, spawn_input(&[("nonexistent-agent", "Do something")])).await?; @@ -211,7 +207,7 @@ async fn spawn_subagent_defaults_to_foreground_without_tasks_capability() -> Tes #[tokio::test] async fn spawn_subagent_preserves_child_result_order_and_ids() -> TestResult { let temp_dir = create_project_with_invocable_agent(); - let mcp = task_client(temp_dir.path()).await?; + let mcp = task_client(temp_dir.path()).await; let parsed = call_complete(&mcp, spawn_input(&[("missing-agent-a", "First task"), ("missing-agent-b", "Second task")])) @@ -228,10 +224,12 @@ async fn spawn_subagent_preserves_child_result_order_and_ids() -> TestResult { #[tokio::test] async fn spawn_subagent_background_returns_mcp_task() -> TestResult { let temp_dir = create_project_with_invocable_agent(); - let mcp = task_client(temp_dir.path()).await?; + let mcp = task_client(temp_dir.path()).await; - let parsed = call_task(&mcp, background_spawn_input(&[("missing-agent", "Do something")])).await?; + let input = background_spawn_input(&[("missing-agent", "Do something")]); + let result = call_task(&mcp, "spawn_subagent", serde_json::to_value(input)?).await?; + let parsed = result.structured_content.ok_or_else(|| test_error("completed task had no structured content"))?; assert_eq!(parsed["results"][0]["status"], "error"); Ok(()) } @@ -239,14 +237,16 @@ async fn spawn_subagent_background_returns_mcp_task() -> TestResult { #[tokio::test] async fn cancelling_background_subagents_cancels_the_batch_without_publishing_a_result() -> TestResult { let temp_dir = create_project_with_blocked_agent(); - let mcp = task_client(temp_dir.path()).await?; - let response = - mcp.raw().call_tool_once(tool_request(background_spawn_input(&[("blocked", "Wait for MCP startup")]))).await?; + let mcp = task_client(temp_dir.path()).await; + let response = mcp + .peer() + .call_tool_once(tool_request("spawn_subagent", background_spawn_input(&[("blocked", "Wait for MCP startup")]))) + .await?; let CallToolResponse::Task(created) = response else { return Err(test_error(format!("expected task response: {response:?}")).into()); }; - mcp.raw().cancel_task(CancelTaskParams::new(created.task.task_id.clone())).await?; + mcp.peer().cancel_task(CancelTaskParams::new(created.task.task_id.clone())).await?; let cancelled = await_terminal_task(&mcp, &created.task.task_id).await?; assert!(matches!(cancelled.payload, TaskPayload::Cancelled)); @@ -254,7 +254,7 @@ async fn cancelling_background_subagents_cancels_the_batch_without_publishing_a_ for _ in 0..3 { tokio::task::yield_now().await; } - let task = mcp.raw().get_task(GetTaskParams::new(created.task.task_id)).await?.task; + let task = mcp.peer().get_task(GetTaskParams::new(created.task.task_id)).await?.task; assert!(matches!(task.payload, TaskPayload::Cancelled)); Ok(()) } @@ -265,24 +265,20 @@ async fn background_instrumentation_finishes_when_the_batch_settles() -> TestRes let project_root = temp_dir.path().to_path_buf(); let finishes = Arc::new(Mutex::new(Vec::new())); let recorded = Arc::clone(&finishes); - let mcp = TestClient::start_with( - move || { - let mut deps = AgentDeps::default().with_agent_registry(test_registry(&project_root)); - deps.observer_factory = Some(Arc::new(RecordingObserverFactory { finishes: recorded })); - SubAgentsMcp::embedded(project_root.clone(), deps) - }, - production_client_info(), - ) - .await?; + let mut deps = AgentDeps::default().with_agent_registry(test_registry(&project_root)); + deps.observer_factory = Some(Arc::new(RecordingObserverFactory { finishes: recorded })); + let mcp = connect(SubAgentsMcp::new(project_root.clone(), deps), cancels_prompts()).await; - let response = - mcp.raw().call_tool_once(tool_request(background_spawn_input(&[("blocked", "Wait for MCP startup")]))).await?; + let response = mcp + .peer() + .call_tool_once(tool_request("spawn_subagent", background_spawn_input(&[("blocked", "Wait for MCP startup")]))) + .await?; let CallToolResponse::Task(created) = response else { return Err(test_error(format!("expected task response: {response:?}")).into()); }; assert!(finishes.lock().unwrap().is_empty(), "span must stay open while the batch is still running"); - mcp.raw().cancel_task(CancelTaskParams::new(created.task.task_id.clone())).await?; + mcp.peer().cancel_task(CancelTaskParams::new(created.task.task_id.clone())).await?; await_terminal_task(&mcp, &created.task.task_id).await?; assert_eq!(finishes.lock().unwrap().as_slice(), [Some("Sub-agent execution cancelled".to_string())]); @@ -292,11 +288,11 @@ async fn background_instrumentation_finishes_when_the_batch_settles() -> TestRes #[tokio::test] async fn spawn_subagent_background_requires_tasks_capability() -> TestResult { let temp_dir = create_project_with_invocable_agent(); - let mcp = TestClient::start_with(|| create_test_server(temp_dir.path()), test_client_info()).await?; + let mcp = connect(create_test_server(temp_dir.path()), taskless()).await; let error = mcp - .raw() - .call_tool_once(tool_request(background_spawn_input(&[("missing-agent", "Do something")]))) + .peer() + .call_tool_once(tool_request("spawn_subagent", background_spawn_input(&[("missing-agent", "Do something")]))) .await .unwrap_err(); @@ -304,93 +300,53 @@ async fn spawn_subagent_background_requires_tasks_capability() -> TestResult { Ok(()) } -async fn task_client(project_root: &Path) -> TestResult> { - TestClient::start_with(|| create_test_server(project_root), production_client_info()).await +async fn task_client(project_root: &Path) -> McpClient { + connect(create_test_server(project_root), cancels_prompts()).await } -async fn await_terminal_task(client: &TestClient, task_id: &str) -> TestResult { - loop { - let task = client.raw().get_task(GetTaskParams::new(task_id.to_string())).await?.task; - if task.status().is_terminal() { - return Ok(task); - } - tokio::task::yield_now().await; - } -} - -fn tool_request(input: SpawnSubAgentsInput) -> CallToolRequestParams { - let arguments = serde_json::to_value(input).unwrap().as_object().unwrap().clone(); - CallToolRequestParams::new("spawn_subagent").with_arguments(arguments) -} - -async fn call_complete(client: &TestClient, input: SpawnSubAgentsInput) -> TestResult { - let response = client.raw().call_tool_once(tool_request(input)).await?; +async fn call_complete(client: &McpClient, input: SpawnSubAgentsInput) -> TestResult { + let response = client.peer().call_tool_once(tool_request("spawn_subagent", input)).await?; let CallToolResponse::Complete(result) = response else { return Err(test_error(format!("expected completed response: {response:?}")).into()); }; result.structured_content.ok_or_else(|| test_error("completed response had no structured content").into()) } -async fn call_task(client: &TestClient, input: SpawnSubAgentsInput) -> TestResult { - let response = client.raw().call_tool_once(tool_request(input)).await?; - let CallToolResponse::Task(created) = response else { - return Err(test_error(format!("expected task response: {response:?}")).into()); - }; - assert_eq!(created.task.status, TaskStatus::Working); - - let detailed = await_terminal_task(client, &created.task.task_id).await?; - let TaskPayload::Completed { result } = detailed.payload else { - return Err(test_error(format!("expected completed task: {detailed:?}")).into()); - }; - let result: CallToolResult = serde_json::from_value(serde_json::Value::Object(result))?; - result.structured_content.ok_or_else(|| test_error("completed task had no structured content").into()) -} - async fn call_subagent_through_manager( project_root: &Path, deps: AgentDeps, input: SpawnSubAgentsInput, ) -> TestResult { - let mut spawn = mcp(project_root) - .with_agent_deps(deps) - .with_builtin_servers() - .from_mcp_config_sources(&[McpConfigSource::Json( - r#"{"servers":{"subagents":{"type":"in-memory","args":[]}}}"#.to_string(), - )])? - .spawn() - .await?; - - let snapshot = spawn.block_until_ready().await.ok_or_else(|| test_error("MCP bootstrap aborted"))?; - let tool = snapshot - .tool_definitions() - .into_iter() - .find(|tool| tool.name.ends_with("spawn_subagent")) - .ok_or_else(|| test_error("spawn_subagent tool missing"))?; - let request = llm::ToolCallRequest { - id: "runtime-registry-test".to_string(), - name: tool.name.clone(), - arguments: serde_json::to_string(&input)?, - }; - - let options = CallToolOptions { timeout: Duration::MAX, meta: None, cancel: CancellationToken::new() }; - let mut events = spawn.handle().call_model_visible(request.name.clone(), &request.arguments, options); + let mcp = ready_catalog(subagents_catalog(project_root, deps)).await; + let arguments = args(serde_json::to_value(&input)?); + let mut events = mcp + .call_tool("subagents__spawn_subagent", arguments, ToolCallOptions::default()) + .expect("the subagents server offers spawn_subagent"); while let Some(event) = events.next().await { match event { - ToolCallEvent::Complete(outcome) => { - let (result, _) = convert_tool_result(&request, outcome, &utils::temp_dir::TempDir::new()) - .map_err(|error| test_error(error.error))?; - return Ok(result.result); - } - ToolCallEvent::TaskComplete { .. } => { + ToolCallEvent::TaskCreated(_) => { return Err(test_error("foreground subagent call unexpectedly returned an MCP Task").into()); } - _ => {} + ToolCallEvent::Done(result) => return Ok(result_text(&result?)), + ToolCallEvent::Progress(_) | ToolCallEvent::TaskStatus(_) => {} } } Err(test_error("MCP manager stopped before returning the tool result").into()) } +fn subagents_catalog(project_root: &Path, deps: AgentDeps) -> McpBuilder { + let config = McpConfig::from_json(r#"{"servers":{"subagents":{"type":"in-memory","args":[]}}}"#).unwrap(); + agent_mcp(project_root, deps, ToolFilter::default()).with_config(config).expect("the subagents config is valid") +} + +fn result_text(result: &CallToolResult) -> String { + match &result.structured_content { + Some(structured) => structured.to_string(), + None => result.content.iter().filter_map(|content| Some(content.as_text()?.text.as_str())).collect(), + } +} + fn create_project_with_invocable_agent() -> TempDir { create_test_files(&[ ( @@ -448,19 +404,8 @@ fn create_project_with_codex_agent() -> TempDir { /// instructions it advertises, which name the agents it will accept. async fn subagent_instructions(project_root: &Path, catalog: AgentCatalog) -> String { let deps = AgentDeps::default().with_agent_registry(catalog.registry().clone()); - let mut spawn = mcp(project_root) - .with_agent_deps(deps) - .with_builtin_servers() - .from_mcp_config_sources(&[McpConfigSource::Json( - r#"{"servers":{"subagents":{"type":"in-memory","args":[]}}}"#.to_string(), - )]) - .expect("Failed to configure subagents MCP") - .spawn() - .await - .expect("Failed to spawn MCP manager"); - - let snapshot = spawn.block_until_ready().await.expect("MCP bootstrap aborted"); - snapshot.model_instructions().get("subagents").expect("Missing subagents instructions").clone() + let mcp = ready_catalog(subagents_catalog(project_root, deps)).await; + mcp.instructions().servers.get("subagents").expect("Missing subagents instructions").clone() } /// The catalog a standalone server would discover from `test_dir`. @@ -475,7 +420,7 @@ fn test_registry(test_dir: &Path) -> aether_core::core::AgentRegistry { fn create_test_server(test_dir: &Path) -> SubAgentsMcp { let deps = AgentDeps::default().with_agent_registry(test_registry(test_dir)); - SubAgentsMcp::embedded(test_dir.to_path_buf(), deps) + SubAgentsMcp::new(test_dir.to_path_buf(), deps) } struct RecordingObserverFactory { diff --git a/crates/mcp-servers/tests/integration/plugins_mcp_skills.rs b/crates/mcp-servers/tests/integration/plugins_mcp_skills.rs index 8e8e64eb3..7c3b92ed4 100644 --- a/crates/mcp-servers/tests/integration/plugins_mcp_skills.rs +++ b/crates/mcp-servers/tests/integration/plugins_mcp_skills.rs @@ -1,5 +1,6 @@ -use crate::common::{TestClient, TestResult, create_test_files, load_skills_input, skills_server}; +use crate::common::{McpClientExt, TestResult, connect, create_test_files, load_skills_input, skills_server}; use mcp_servers::skills::tools::ListSkillsInput; +use mcp_utils::client::ClientOptions; use serde::{Deserialize, Serialize}; use std::path::PathBuf; use utils::MarkdownFile; @@ -55,7 +56,7 @@ async fn test_load_skills_tool() -> TestResult { ]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("get_skills", load_skills_input(&[("skill-1", None), ("skill-2", None), ("skill-3", None)])).await?; @@ -94,7 +95,7 @@ async fn test_list_skills_only_returns_agent_invocable_entries() -> TestResult { ]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("list_skills", ListSkillsInput::default()).await?; @@ -122,7 +123,7 @@ async fn test_load_skills_with_missing() -> TestResult { ]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp .call("get_skills", load_skills_input(&[("skill-1", None), ("nonexistent-skill", None), ("skill-2", None)])) @@ -160,7 +161,7 @@ async fn test_get_skills_rejects_non_agent_invocable_prompts() -> TestResult { ]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp .call("get_skills", load_skills_input(&[("allowed", None), ("user-only", None), ("rule-only", None)])) @@ -194,7 +195,7 @@ async fn test_load_auxiliary_file() -> TestResult { ]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("get_skills", load_skills_input(&[("test-skill", None)])).await?; let file = &parsed["files"][0]; @@ -217,7 +218,7 @@ async fn test_reject_traversal() -> TestResult { let test_files = vec![("skills/test-skill/SKILL.md", "---\ndescription: Test\nagent-invocable: true\n---\n# Test")]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("get_skills", load_skills_input(&[("test-skill", Some("../other-skill/SKILL.md"))])).await?; let file = &parsed["files"][0]; @@ -231,7 +232,7 @@ async fn test_reject_absolute_path() -> TestResult { let test_files = vec![("skills/test-skill/SKILL.md", "---\ndescription: Test\nagent-invocable: true\n---\n# Test")]; let temp_dir = create_test_files(&test_files); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("get_skills", load_skills_input(&[("test-skill", Some("/etc/passwd"))])).await?; let file = &parsed["files"][0]; @@ -243,9 +244,9 @@ async fn test_reject_absolute_path() -> TestResult { #[tokio::test] async fn list_skills_input_schema_has_properties_object() -> TestResult { let temp_dir = create_test_files(&[]); - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; - let tools = mcp.raw().peer().list_all_tools().await?; + let tools = mcp.list_tools().await?; let tool = tools.into_iter().find(|tool| tool.name.as_ref() == "list_skills").expect("list_skills tool present"); let schema = serde_json::Value::Object((*tool.input_schema).clone()); diff --git a/crates/mcp-servers/tests/integration/review_mcp.rs b/crates/mcp-servers/tests/integration/review_mcp.rs index 2c62ee357..09246a276 100644 --- a/crates/mcp-servers/tests/integration/review_mcp.rs +++ b/crates/mcp-servers/tests/integration/review_mcp.rs @@ -1,7 +1,8 @@ -use crate::common::{TestClient, TestResult, scripted_mcp_client, silent_mcp_client, test_error}; +use crate::common::{McpClientExt, TestResult, cancels_prompts, connect, scripted, test_error}; use axum::{Router, extract::Query, http::header::CONTENT_TYPE, response::Html, routing::get}; use mcp_servers::review::ReviewMcp; -use mcp_utils::{client::McpClient, testing::ElicitationScript}; +use mcp_utils::client::{ClientOptions, McpClient}; +use mcp_utils::testing::ElicitationScript; use reqwest::Url; use rmcp::model::{ CallToolRequest, CallToolRequestParams, CallToolResponse, ClientRequest, ElicitRequestParams, ElicitResult, @@ -162,7 +163,7 @@ async fn cancelling_the_tool_call_tears_down_the_review_server() -> TestResult { let second_round = response_request(args, ElicitResult::new(ElicitationAction::Accept)).with_request_state(review.token.clone()); let request = ClientRequest::CallToolRequest(CallToolRequest::new(second_round)); - let handle = test.client.raw().send_cancellable_request(request, PeerRequestOptions::no_options()).await?; + let handle = test.client.peer().send_cancellable_request(request, PeerRequestOptions::no_options()).await?; handle.cancel(None).await?; while reqwest::get(&review.url).await.is_ok() { @@ -266,11 +267,13 @@ async fn explicit_approval_returns_approved_without_feedback() -> TestResult { async fn clients_without_elicitation_support_are_rejected() -> TestResult { let root = TempDir::new()?; fs::write(root.path().join("valid.md"), "# Valid")?; - let mcp = TestClient::start(|| review_mcp_at(root.path())).await?; + let mcp = connect(review_mcp_at(root.path()), ClientOptions::default()).await; let result = mcp.call_raw("review_artifact", markdown_file("valid.md")).await?; assert_eq!(result.is_error, Some(true)); + let text = result.content.first().and_then(|content| content.as_text()).expect("error text"); + assert!(text.text.contains("cannot prompt the user"), "unexpected error: {}", text.text); Ok(()) } @@ -286,7 +289,7 @@ struct ReviewTestBuilder { struct ReviewTest { root: TempDir, - client: TestClient, + client: McpClient, script: Option, } @@ -330,14 +333,14 @@ impl ReviewTestBuilder { fs::write(path, content)?; } - let (client, script) = self.response.map_or_else( - || (silent_mcp_client("review-test-server"), None), - |response| { - let (client, script) = scripted_mcp_client("review-test-server", response); - (client, Some(script)) - }, - ); - let client = TestClient::start_with(|| review_mcp_at(root.path()), client).await?; + let (options, script) = match self.response { + Some(response) => { + let (options, script) = scripted(response); + (options, Some(script)) + } + None => (cancels_prompts(), None), + }; + let client = connect(review_mcp_at(root.path()), options).await; Ok(ReviewTest { root, client, script }) } } @@ -358,12 +361,12 @@ impl ReviewTest { /// The first MRTR round of a Markdown review. async fn request_form(&self, args: Value) -> TestResult { - FormReview::from_response(self.client.raw().call_tool_once(tool_request(args)).await?) + FormReview::from_response(self.client.peer().call_tool_once(tool_request(args)).await?) } /// The first MRTR round of an HTML review. async fn request_url(&self, args: Value) -> TestResult { - UrlReview::from_response(self.client.raw().call_tool_once(tool_request(args)).await?) + UrlReview::from_response(self.client.peer().call_tool_once(tool_request(args)).await?) } /// The second MRTR round: the client echoes the arguments, the review token @@ -373,7 +376,7 @@ impl ReviewTest { if let Some(token) = token { request = request.with_request_state(token.to_string()); } - structured_output(self.client.raw().call_tool_once(request).await?) + structured_output(self.client.peer().call_tool_once(request).await?) } /// A full review driven by the scripted elicitation response. @@ -465,7 +468,7 @@ impl Drop for UpstreamApp { } fn review_mcp_at(root: &Path) -> ReviewMcp { - ReviewMcp::from_args_with_base_dir(Vec::new(), root).expect("empty review MCP arguments are valid") + ReviewMcp::new(root.to_path_buf()) } fn accepted_response(content: Value) -> ElicitResult { @@ -507,9 +510,8 @@ fn with_title(mut args: Value, title: Option<&str>) -> Value { args } -fn tool_request(args: Value) -> CallToolRequestParams { - let Value::Object(arguments) = args else { panic!("tool arguments") }; - CallToolRequestParams::new("review_artifact").with_arguments(arguments) +fn tool_request(arguments: Value) -> CallToolRequestParams { + crate::common::tool_request("review_artifact", arguments) } fn response_request(args: Value, result: ElicitResult) -> CallToolRequestParams { diff --git a/crates/mcp-servers/tests/integration/skills_self_improvement.rs b/crates/mcp-servers/tests/integration/skills_self_improvement.rs index 3aa8477b6..919639b44 100644 --- a/crates/mcp-servers/tests/integration/skills_self_improvement.rs +++ b/crates/mcp-servers/tests/integration/skills_self_improvement.rs @@ -1,6 +1,7 @@ -use crate::common::{TestClient, TestResult, load_skills_input, skills_server}; +use crate::common::{McpClientExt, TestResult, connect, load_skills_input, skills_server}; use mcp_servers::skills::SkillsMcp; use mcp_servers::skills::tools::ListSkillsInput; +use mcp_utils::client::ClientOptions; use rmcp::ServerHandler; use tempfile::TempDir; @@ -23,7 +24,7 @@ async fn test_instructions_reference_list_skills_and_do_not_embed_catalog_entrie std::fs::write(human_dir.join("SKILL.md"), "---\ndescription: Human skill\nagent-invocable: true\n---\nContent.\n") .unwrap(); - let server = SkillsMcp::new(&[skills_dir]); + let server = SkillsMcp::new(temp_dir.path(), &[skills_dir]); let info = server.get_info(); let instructions = info.instructions.unwrap(); @@ -45,7 +46,7 @@ async fn test_full_lifecycle() -> TestResult { "---\ndescription: Curated skill\nagent-invocable: true\n---\n# Curated\n\nHand-written skill.", )?; - let mcp = TestClient::start(|| skills_server(temp_dir.path())).await?; + let mcp = connect(skills_server(temp_dir.path()), ClientOptions::default()).await; let parsed = mcp.call("list_skills", ListSkillsInput::default()).await?; let skills = parsed["skills"].as_array().unwrap(); diff --git a/crates/mcp-servers/tests/integration/stdio_transport.rs b/crates/mcp-servers/tests/integration/stdio_transport.rs index f4f01bfce..cdb3924cf 100644 --- a/crates/mcp-servers/tests/integration/stdio_transport.rs +++ b/crates/mcp-servers/tests/integration/stdio_transport.rs @@ -1,36 +1,26 @@ -#![cfg(feature = "stdio")] - -use rmcp::ServiceExt; -use rmcp::model::CallToolRequestParams; -use rmcp::transport::TokioChildProcess; -use std::process::Stdio; -use tokio::process::Command; - -fn stdio_binary() -> Command { - let mut cmd = Command::new(env!("CARGO_BIN_EXE_mcp-servers-stdio")); - cmd.stdin(Stdio::piped()).stdout(Stdio::piped()); - cmd +#![cfg(feature = "subagents")] + +use crate::common::call_tool; +use mcp_utils::client::{ClientError, ClientOptions, McpClient, Transport}; +use serde_json::json; +use std::collections::HashMap; + +async fn connect(args: &[&str], env: HashMap) -> Result { + let command = env!("CARGO_BIN_EXE_mcp-servers-stdio").to_string(); + let args = args.iter().map(ToString::to_string).collect(); + McpClient::connect("stdio", Transport::Stdio { command, args, env, cwd: None }, &ClientOptions::default()).await } fn tool_names(tools: &[rmcp::model::Tool]) -> Vec<&str> { tools.iter().map(|t| t.name.as_ref()).collect() } -fn extract_text(content: &rmcp::model::ContentBlock) -> &str { - content.as_text().expect("expected text content").text.as_str() -} - async fn connect_and_list_tools(server: &str, extra_args: &[&str]) -> Vec { let aether_home = tempfile::tempdir().expect("create temp aether home"); - let mut cmd = stdio_binary(); - cmd.env("AETHER_HOME", aether_home.path().join(".aether")); - cmd.arg("--server").arg(server); - for arg in extra_args { - cmd.arg(arg); - } - let transport = TokioChildProcess::new(cmd).expect("spawn stdio server"); - let client = ().serve(transport).await.expect("connect to server"); - client.peer().list_all_tools().await.expect("list tools") + let env = HashMap::from([("AETHER_HOME".to_string(), aether_home.path().join(".aether").display().to_string())]); + let args = [&["--server", server], extra_args].concat(); + let client = connect(&args, env).await.expect("connect to server"); + client.list_tools().await.expect("list tools") } #[tokio::test] @@ -50,47 +40,18 @@ async fn tasks_server_lists_tools_over_stdio() { async fn tasks_server_create_and_get_task_over_stdio() { let tmp = tempfile::tempdir().expect("create temp dir"); - let mut cmd = stdio_binary(); - cmd.arg("--server").arg("tasks"); - cmd.arg("--").arg("--dir").arg(tmp.path()); - - let transport = TokioChildProcess::new(cmd).expect("spawn stdio server"); - let client = ().serve(transport).await.expect("connect to server"); - - // This test exercises call_tool, so it doesn't use connect_and_list_tools. - // Create a task - let create_result = client - .peer() - .call_tool( - rmcp::model::CallToolRequestParams::new("task_create").with_arguments( - serde_json::json!({ - "title": "Test task", - "description": "A test task created over stdio" - }) - .as_object() - .unwrap() - .clone(), - ), - ) - .await - .expect("call task_create"); - - let text = extract_text(create_result.content.first().expect("response has content")); - let created: serde_json::Value = serde_json::from_str(text).expect("parse JSON response"); + let dir = tmp.path().to_str().unwrap(); + let client = connect(&["--server", "tasks", "--", "--dir", dir], HashMap::new()).await.expect("connect to server"); + + let created = call_tool( + &client, + "task_create", + json!({ "title": "Test task", "description": "A test task created over stdio" }), + ) + .await; let task_id = created["task"]["id"].as_str().expect("task has id"); - // Get the task back - let get_result = client - .peer() - .call_tool( - CallToolRequestParams::new("task_get") - .with_arguments(serde_json::json!({ "id": task_id }).as_object().unwrap().clone()), - ) - .await - .expect("call task_get"); - - let get_text = extract_text(get_result.content.first().expect("response has content")); - let fetched: serde_json::Value = serde_json::from_str(get_text).expect("parse JSON response"); + let fetched = call_tool(&client, "task_get", json!({ "id": task_id })).await; assert_eq!(fetched["task"]["title"].as_str(), Some("Test task")); } @@ -127,11 +88,7 @@ async fn coding_server_accepts_rules_dir_over_stdio() { #[tokio::test] async fn lsp_server_is_not_available_over_stdio() { - let mut cmd = stdio_binary(); - cmd.arg("--server").arg("lsp"); - - let transport = TokioChildProcess::new(cmd).expect("spawn stdio server"); - let result = ().serve(transport).await; + let result = connect(&["--server", "lsp"], HashMap::new()).await; assert!(result.is_err(), "expected lsp server to be unavailable"); } @@ -184,11 +141,7 @@ async fn review_server_lists_tools_over_stdio() { #[tokio::test] async fn unknown_server_exits_with_error() { - let mut cmd = stdio_binary(); - cmd.arg("--server").arg("nonexistent"); - - let transport = TokioChildProcess::new(cmd).expect("spawn process"); - let result = ().serve(transport).await; + let result = connect(&["--server", "nonexistent"], HashMap::new()).await; assert!(result.is_err(), "expected error for unknown server name"); } diff --git a/crates/mcp-servers/tests/integration/test_bash.rs b/crates/mcp-servers/tests/integration/test_bash.rs index febe56f75..9588eb7d5 100644 --- a/crates/mcp-servers/tests/integration/test_bash.rs +++ b/crates/mcp-servers/tests/integration/test_bash.rs @@ -1,11 +1,13 @@ use mcp_servers::CodingMcp; use mcp_servers::coding::error::BashError; use mcp_servers::coding::tools::bash::{BashEnvironment, BashInput, execute_command}; -use rmcp::model::{CallToolRequestParams, CallToolResponse, TaskPayload, TaskStatus}; +use mcp_utils::testing::args; +use rmcp::model::{CallToolRequestParams, CallToolResponse, CancelTaskParams, TaskStatus}; +use serde_json::json; use std::fs::canonicalize; use std::process::{Command, Stdio}; -use super::common::{TestClient, TestResult, production_client_info, test_client_info}; +use super::common::{TestResult, await_terminal_task, call_task, cancels_prompts, connect, taskless}; #[tokio::test] async fn bash_commands_receive_eof_instead_of_inherited_stdin() { @@ -113,37 +115,14 @@ async fn command_runs_in_working_directory() -> Result<(), Box TestResult { - let client = TestClient::start_with(CodingMcp::new, production_client_info()).await?; - let response = client - .raw() - .call_tool_once( - CallToolRequestParams::new("bash").with_arguments( - serde_json::json!({ - "command": "echo stdout; echo stderr >&2; exit 7", - "runInBackground": true - }) - .as_object() - .unwrap() - .clone(), - ), - ) - .await?; - let CallToolResponse::Task(created) = response else { - panic!("expected task response: {response:?}"); - }; - assert_eq!(created.task.status, TaskStatus::Working); - - let detailed = loop { - let result = client.raw().get_task(rmcp::model::GetTaskParams::new(created.task.task_id.clone())).await?; - if result.task.status().is_terminal() { - break result.task; - } - tokio::task::yield_now().await; - }; - let TaskPayload::Completed { result } = detailed.payload else { - panic!("expected completed task"); - }; - let result: rmcp::model::CallToolResult = serde_json::from_value(serde_json::Value::Object(result))?; + let client = connect(CodingMcp::new(std::env::temp_dir()), cancels_prompts()).await; + let result = call_task( + &client, + "bash", + json!({ "command": "echo stdout; echo stderr >&2; exit 7", "runInBackground": true }), + ) + .await?; + assert_eq!(result.is_error, Some(false)); let structured = result.structured_content.as_ref().unwrap(); let output = structured["output"].as_str().unwrap(); @@ -153,14 +132,31 @@ async fn background_bash_uses_tasks_and_returns_structured_output() -> TestResul Ok(()) } +#[tokio::test] +async fn cancelling_a_background_bash_task_stops_it() -> TestResult { + let client = connect(CodingMcp::new(std::env::temp_dir()), cancels_prompts()).await; + let request = CallToolRequestParams::new("bash") + .with_arguments(args(json!({ "command": "sleep 60", "runInBackground": true }))); + let CallToolResponse::Task(created) = client.peer().call_tool_once(request).await? else { + panic!("background bash creates a task"); + }; + + client.peer().cancel_task(CancelTaskParams::new(created.task.task_id.clone())).await?; + + let task = await_terminal_task(&client, &created.task.task_id).await?; + assert_eq!(task.status(), TaskStatus::Cancelled); + Ok(()) +} + #[tokio::test] async fn background_bash_requires_tasks_capability() -> TestResult { - let client = TestClient::start_with(CodingMcp::new, test_client_info()).await?; + let client = connect(CodingMcp::new(std::env::temp_dir()), taskless()).await; let error = client - .raw() - .call_tool_once(CallToolRequestParams::new("bash").with_arguments( - serde_json::json!({ "command": "echo never", "runInBackground": true }).as_object().unwrap().clone(), - )) + .peer() + .call_tool_once( + CallToolRequestParams::new("bash") + .with_arguments(args(json!({ "command": "echo never", "runInBackground": true }))), + ) .await .unwrap_err(); assert!( diff --git a/crates/mcp-servers/tests/integration/test_read_before_edit_safety.rs b/crates/mcp-servers/tests/integration/test_read_before_edit_safety.rs index c6840657b..d8d1d82d5 100644 --- a/crates/mcp-servers/tests/integration/test_read_before_edit_safety.rs +++ b/crates/mcp-servers/tests/integration/test_read_before_edit_safety.rs @@ -18,7 +18,7 @@ async fn test_edit_file_without_read_fails() { let test_file = temp_dir.path().join("test.txt"); fs::write(&test_file, "original content").unwrap(); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // Try to edit without reading first let edit_args = @@ -39,7 +39,7 @@ async fn test_edit_file_after_read_succeeds() { let test_file = temp_dir.path().join("test.txt"); fs::write(&test_file, "original content").unwrap(); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // First read the file let read_args = ReadFileArgs { file_path: test_file.to_string_lossy().to_string(), offset: None, limit: None }; @@ -65,7 +65,7 @@ async fn test_write_existing_file_without_read_fails() { let test_file = temp_dir.path().join("test.txt"); fs::write(&test_file, "original content").unwrap(); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // Try to write without reading first let write_args = @@ -86,7 +86,7 @@ async fn test_write_existing_file_after_read_succeeds() { let test_file = temp_dir.path().join("test.txt"); fs::write(&test_file, "original content").unwrap(); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // First read the file let read_args = ReadFileArgs { file_path: test_file.to_string_lossy().to_string(), offset: None, limit: None }; @@ -111,7 +111,7 @@ async fn test_write_new_file_without_read_succeeds() { let temp_dir = TempDir::new().unwrap(); let test_file = temp_dir.path().join("new_file.txt"); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // Write to a file that doesn't exist - should succeed let write_args = @@ -135,7 +135,7 @@ async fn test_multiple_files_tracked_independently() { fs::write(&file1, "content 1").unwrap(); fs::write(&file2, "content 2").unwrap(); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // Read file1 let read_args = ReadFileArgs { file_path: file1.to_string_lossy().to_string(), offset: None, limit: None }; @@ -156,7 +156,7 @@ async fn test_failed_read_doesnt_track_file() { let temp_dir = TempDir::new().unwrap(); let nonexistent_file = temp_dir.path().join("doesnt_exist.txt"); - let mcp = CodingMcp::new(); + let mcp = CodingMcp::new(temp_dir.path()); // Try to read a file that doesn't exist let read_args = diff --git a/crates/mcp-servers/tests/integration/test_read_file_rule_meta.rs b/crates/mcp-servers/tests/integration/test_read_file_rule_meta.rs index aa08340da..71e9b48cb 100644 --- a/crates/mcp-servers/tests/integration/test_read_file_rule_meta.rs +++ b/crates/mcp-servers/tests/integration/test_read_file_rule_meta.rs @@ -21,7 +21,7 @@ async fn test_read_file_meta_includes_matched_rule_names_from_configured_rules_d let rs_file = src_dir.join("main.rs"); fs::write(&rs_file, "fn main() {}\n").unwrap(); - let mcp = CodingMcp::new().with_rules_dirs(vec![rules_dir]).with_root_dir(root.to_path_buf()); + let mcp = CodingMcp::new(root).with_rules_dirs(vec![rules_dir]); let result = mcp .test_read_file(ReadFileArgs { file_path: rs_file.to_string_lossy().to_string(), offset: None, limit: None }) @@ -59,7 +59,7 @@ async fn test_read_file_does_not_auto_load_rules_without_rules_dirs() { let rs_file = src_dir.join("main.rs"); fs::write(&rs_file, "fn main() {}\n").unwrap(); - let mcp = CodingMcp::new().with_root_dir(root.to_path_buf()); + let mcp = CodingMcp::new(root); let result = mcp .test_read_file(ReadFileArgs { file_path: rs_file.to_string_lossy().to_string(), offset: None, limit: None }) diff --git a/crates/mcp-servers/tests/lsp_symbol_errors_e2e.rs b/crates/mcp-servers/tests/lsp_symbol_errors_e2e.rs index c4f105f80..da18ac1b3 100644 --- a/crates/mcp-servers/tests/lsp_symbol_errors_e2e.rs +++ b/crates/mcp-servers/tests/lsp_symbol_errors_e2e.rs @@ -3,7 +3,8 @@ mod common; use aether_lspd::LanguageId; use aether_lspd::testing::{CargoProject, TestDaemon, TestProject, configure_fake_server}; -use common::{call_tool_error, connect_lsp}; +use common::{call_tool_error, coding_with_lsp, connect}; +use mcp_utils::client::ClientOptions; use std::time::Duration; #[tokio::test] @@ -13,7 +14,7 @@ async fn symbol_lookup_propagates_language_server_errors() { project.add_file("src/lib.rs", "pub fn example_fn() {}\n").expect("add source file"); let daemon = TestDaemon::spawn(project.root(), LanguageId::Rust, Duration::from_secs(120)).await.expect("start daemon"); - let (server_handle, client) = connect_lsp(&project).await; + let client = connect(coding_with_lsp(&project), ClientOptions::default()).await; let error = call_tool_error( &client, @@ -23,7 +24,6 @@ async fn symbol_lookup_propagates_language_server_errors() { .await; assert!(error.contains("fake server rejected request"), "unexpected error: {error}"); - drop(client); - drop(server_handle); + client.close().await; daemon.shutdown().expect("stop daemon"); } diff --git a/crates/mcp-servers/tests/lsp_symbol_recovery_e2e.rs b/crates/mcp-servers/tests/lsp_symbol_recovery_e2e.rs index e4f88d3cb..9163fd6d2 100644 --- a/crates/mcp-servers/tests/lsp_symbol_recovery_e2e.rs +++ b/crates/mcp-servers/tests/lsp_symbol_recovery_e2e.rs @@ -3,9 +3,9 @@ mod common; use aether_lspd::testing::{TestDaemon, configure_fake_server}; use aether_lspd::{LanguageId, socket_path}; -use common::{call_tool, call_tool_error, test_client_info}; +use common::{call_tool, call_tool_error, connect}; use mcp_servers::coding::CodingMcp; -use mcp_utils::testing::connect; +use mcp_utils::client::ClientOptions; use std::time::Duration; use tempfile::tempdir; @@ -22,8 +22,8 @@ async fn symbol_lookup_recovers_after_language_server_timeout() { let daemon = TestDaemon::spawn(root.path(), LanguageId::TypeScript, Duration::from_secs(1)) .await .expect("Failed to spawn test daemon"); - let server = CodingMcp::new().with_lsp(root.path().to_path_buf()); - let (server_handle, client) = connect(server, test_client_info()).await.expect("Failed to connect"); + let server = CodingMcp::new(root.path()).with_lsp(); + let client = connect(server, ClientOptions::default()).await; let lookup = |operation: &str| serde_json::json!({ "operation": operation, "file_path": "example.ts", "symbol": "example_fn", "line": 1 }); let error = call_tool_error(&client, "lsp_symbol", lookup("definition")).await; @@ -33,8 +33,7 @@ async fn symbol_lookup_recovers_after_language_server_timeout() { let hover = result["hoverContents"].as_str().unwrap_or_default(); assert!(hover.contains("export function example_fn"), "{result}"); - drop(client); - drop(server_handle); + client.close().await; daemon.shutdown().expect("Failed to shut down test daemon"); assert!(!socket_path(root.path(), LanguageId::TypeScript).exists()); } diff --git a/crates/mcp-utils/Cargo.toml b/crates/mcp-utils/Cargo.toml index 223f11e11..e54294c4d 100644 --- a/crates/mcp-utils/Cargo.toml +++ b/crates/mcp-utils/Cargo.toml @@ -25,40 +25,60 @@ workspace = true [features] default = ["client"] -testing = ["client"] +testing = ["client", "dep:axum", "rmcp/transport-streamable-http-server"] client = [ "dep:aether-auth", "aether-auth/mcp", "dep:async-stream", - "dep:llm", + "dep:chrono", "dep:reqwest", + "dep:serde", + "dep:serde_json", + "dep:tokio-stream", "rmcp/transport-streamable-http-client-reqwest", "rmcp/transport-child-process", ] [dependencies] -chrono = { workspace = true } futures = { workspace = true } -serde = { workspace = true } +rmcp = { workspace = true, features = ["client", "elicitation", "server", "transport-async-rw", "transport-io"] } +thiserror = { workspace = true } tokio = { workspace = true, features = ["rt", "macros", "sync", "net", "time"] } -tokio-util = { workspace = true } -rmcp = { workspace = true, features = ["client", "elicitation", "server", "transport-async-rw"] } +tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } -utils = { package = "aether-utils", path = "../utils", version = "0.2.19" } -schemars = { workspace = true } +uuid = { workspace = true } # client-only deps +axum = { workspace = true, optional = true } aether-auth = { path = "../aether-auth", version = "0.1.17", optional = true } async-stream = { workspace = true, optional = true } -llm = { package = "aether-llm", path = "../llm", optional = true, version = "0.9.13" } -serde_json = { workspace = true } +chrono = { workspace = true, optional = true } reqwest = { workspace = true, optional = true } -thiserror = { workspace = true } -uuid = { workspace = true } +serde = { workspace = true, optional = true } +serde_json = { workspace = true, optional = true } +tokio-stream = { workspace = true, features = ["sync"], optional = true } + +[[test]] +name = "client" +required-features = ["testing"] + +[[test]] +name = "oauth" +required-features = ["testing"] + +[[test]] +name = "gateway" +required-features = ["testing"] + +[[test]] +name = "server" +required-features = ["testing"] + +[[test]] +name = "http_headers" +required-features = ["client"] [dev-dependencies] axum = { workspace = true } -aether-doctest = { path = "../doctest" } -rmcp = { workspace = true, features = ["macros"] } -tempfile = { workspace = true } -tracing-subscriber = { workspace = true } +rmcp = { workspace = true, features = ["macros", "transport-streamable-http-server"] } +tokio = { workspace = true, features = ["rt-multi-thread", "test-util"] } diff --git a/crates/mcp-utils/README.md b/crates/mcp-utils/README.md index 29788ba1f..fbdba5928 100644 --- a/crates/mcp-utils/README.md +++ b/crates/mcp-utils/README.md @@ -1,13 +1,14 @@ # aether-mcp-utils -Utilities for the [Model Context Protocol](https://modelcontextprotocol.io/) (MCP), providing transport, status tracking, and client management for MCP servers. +[Model Context Protocol](https://modelcontextprotocol.io/) (MCP) building blocks on top of +[rmcp](https://docs.rs/rmcp): connect to a server, host a server, or aggregate many servers +behind one gateway. ## Table of Contents -- [Key Types](#key-types) - [Feature Flags](#feature-flags) - [License](#license) @@ -17,7 +18,8 @@ Utilities for the [Model Context Protocol](https://modelcontextprotocol.io/) (MC | Feature | Description | Default | |---------|-------------|---------| -| `client` | MCP client with OAuth, server management, and deferred tool routing | yes | +| `client` | `client` and `gateway`, including OAuth. Without it, only server hosting remains. | yes | +| `testing` | `testing`: an in-memory `FakeMcpServer`, an OAuth-protected `FakeAuthServer`, and a scripted elicitation host | no | ## License diff --git a/crates/mcp-utils/src/client/call_tool.rs b/crates/mcp-utils/src/client/call_tool.rs deleted file mode 100644 index 4439cf29c..000000000 --- a/crates/mcp-utils/src/client/call_tool.rs +++ /dev/null @@ -1,221 +0,0 @@ -use crate::client::McpClient; -use crate::client::elicitation::{ElicitInputsError, elicit_inputs}; -use crate::client::mrtr::{AbortReason, MrtrAction, MrtrState}; -use crate::client::task::{TaskDriver, TaskErrorReason}; -use async_stream::stream; -use futures::Stream; -use futures::StreamExt; -use futures::future::{Either, select}; -use rmcp::RoleClient; -use rmcp::model::{ - CallToolRequestParams, CallToolResponse, CallToolResult, ClientRequest, CreateTaskResult, InputRequests, - InputRequiredResult, InputResponses, ProgressNotificationParam, Request, RequestMetaObject, ServerResult, Task, -}; -use rmcp::service::{PeerRequestOptions, RequestHandle, RunningService, ServiceError}; -use std::pin::pin; -use std::sync::Arc; -use std::time::Duration; -use thiserror::Error; -use tokio::time::sleep; -use tokio_util::sync::CancellationToken; - -#[derive(Debug, Default)] -pub struct CallToolOptions { - pub timeout: Duration, - pub meta: Option, - pub cancel: CancellationToken, -} - -#[derive(Debug)] -pub enum ToolCallEvent { - Progress(ProgressNotificationParam), - TaskCreated(CreateTaskResult), - TaskStatus(Task), - Complete(Result), - TaskComplete { task: Task, result: Result }, - Cancelled { task_id: Option }, -} - -#[derive(Debug, Error)] -pub enum CallToolError { - #[error("Failed to send tool request: {0}")] - Send(#[source] ServiceError), - #[error("Tool execution failed: {0}")] - Call(#[source] ServiceError), - #[error("Server '{server}' requested an input kind this client does not support (sampling or roots)")] - UnsupportedInput { server: String }, - #[error("Server '{server}' failed to serialize an elicitation response: {source}")] - SerializeResponse { - server: String, - #[source] - source: serde_json::Error, - }, - #[error("{}", reason.message(server, *timeout))] - Aborted { server: String, reason: AbortReason, timeout: Duration }, - #[error("Task '{task_id}' from server '{server}' {reason}")] - Task { - server: String, - task_id: String, - #[source] - reason: Box, - }, - #[error("Server '{server}' returned a tool call response kind this client does not support")] - UnsupportedResponse { server: String }, - #[error("{message}")] - Unavailable { message: String }, -} - -pub fn call_tool( - client: Arc>, - mut params: CallToolRequestParams, - options: CallToolOptions, -) -> impl Stream { - stream! { - let server_name = client.service().server_name().to_string(); - let mut mrtr_state = MrtrState::new(options.timeout); - - loop { - let request = ClientRequest::CallToolRequest(Request::new(params.clone())); - let send = pin!(client.send_cancellable_request(request, peer_request_options(&options))); - let handle = match select(send, pin!(options.cancel.cancelled())).await { - Either::Left((Ok(handle), _)) => handle, - Either::Left((Err(e), _)) => { - yield ToolCallEvent::Complete(Err(CallToolError::Send(e))); - return; - } - Either::Right(((), _)) => { - yield ToolCallEvent::Cancelled { task_id: None }; - return; - } - }; - - let mut progress = client.service().progress_dispatcher.subscribe(handle.progress_token.clone()).await; - let mut response_or_cancel = pin!(await_response_or_cancel(handle, &options.cancel)); - let response = loop { - match select(progress.next(), response_or_cancel.as_mut()).await { - Either::Left((Some(progress), _)) => yield ToolCallEvent::Progress(progress), - Either::Left((None, response_or_cancel)) => break response_or_cancel.await, - Either::Right((response, _)) => break response, - } - }; - let Some(response) = response else { - yield ToolCallEvent::Cancelled { task_id: None }; - return; - }; - - match response { - Ok(CallToolResponse::Complete(result)) => { - yield ToolCallEvent::Complete(Ok(result)); - return; - } - Ok(CallToolResponse::InputRequired(input_required)) => { - match mrtr_state.tick(input_required) { - MrtrAction::Poll { backoff, request_state } => { - let backoff = pin!(sleep(backoff)); - if let Either::Right(((), _)) = select(backoff, pin!(options.cancel.cancelled())).await { - yield ToolCallEvent::Cancelled { task_id: None }; - return; - } - params.input_responses = None; - params.request_state = Some(request_state); - } - MrtrAction::Elicit { input_requests, request_state } => { - let elicit = pin!(elicit_input(client.service(), &mut mrtr_state, &server_name, input_requests)); - match select(elicit, pin!(options.cancel.cancelled())).await { - Either::Left((Ok(responses), _)) => { - params.input_responses = Some(responses); - params.request_state = request_state; - } - Either::Left((Err(e), _)) => { - yield ToolCallEvent::Complete(Err(e)); - return; - } - Either::Right(((), _)) => { - yield ToolCallEvent::Cancelled { task_id: None }; - return; - } - } - } - MrtrAction::Abort(reason) => { - yield ToolCallEvent::Complete(Err(CallToolError::Aborted { - server: server_name, - reason, - timeout: options.timeout, - })); - return; - } - } - } - Ok(CallToolResponse::Task(task)) => { - let driver = TaskDriver::new(&server_name, client.as_ref(), options.timeout, options.cancel.clone()); - let mut events = Box::pin(driver.stream(task, progress)); - while let Some(event) = events.next().await { - yield event; - } - return; - } - Ok(_) => { - yield ToolCallEvent::Complete(Err(CallToolError::UnsupportedResponse { server: server_name })); - return; - } - Err(e) => { - yield ToolCallEvent::Complete(Err(CallToolError::Call(e))); - return; - } - } - } - } -} - -fn peer_request_options(options: &CallToolOptions) -> PeerRequestOptions { - let request_options = PeerRequestOptions::with_timeout(options.timeout); - match &options.meta { - Some(meta) => request_options.with_meta(meta.clone()), - None => request_options, - } -} - -async fn await_response_or_cancel( - handle: RequestHandle, - cancel: &CancellationToken, -) -> Option> { - let response = pin!(await_tool_response(handle)); - match select(response, pin!(cancel.cancelled())).await { - Either::Left((response, _)) => Some(response), - Either::Right(((), _)) => None, - } -} - -async fn await_tool_response(handle: RequestHandle) -> Result { - match handle.await_response().await? { - ServerResult::CallToolResult(result) => Ok(CallToolResponse::Complete(result)), - ServerResult::InputRequiredResult(result) => Ok(CallToolResponse::InputRequired(result)), - ServerResult::CreateTaskResult(result) => Ok(CallToolResponse::Task(result)), - ServerResult::CustomResult(result) - if result.0.get("resultType").and_then(serde_json::Value::as_str) == Some("input_required") - && result.0.get("inputRequests").is_none() - && result.0.get("requestState").is_none() => - { - Ok(CallToolResponse::InputRequired(InputRequiredResult::new(None, None))) - } - _ => Err(ServiceError::UnexpectedResponse), - } -} - -async fn elicit_input( - client: &McpClient, - mrtr_state: &mut MrtrState, - server_name: &str, - requests: InputRequests, -) -> Result { - let (responses, results) = elicit_inputs(client, requests).await.map_err(|error| match error { - ElicitInputsError::UnsupportedInput => CallToolError::UnsupportedInput { server: server_name.to_string() }, - ElicitInputsError::Serialize(source) => { - CallToolError::SerializeResponse { server: server_name.to_string(), source } - } - })?; - for result in &results { - mrtr_state.record_response(result); - } - Ok(responses) -} diff --git a/crates/mcp-utils/src/client/config.rs b/crates/mcp-utils/src/client/config.rs deleted file mode 100644 index 2e6b621c5..000000000 --- a/crates/mcp-utils/src/client/config.rs +++ /dev/null @@ -1,804 +0,0 @@ -use aether_auth::OAuthClientRegistration; -use reqwest::header::{AUTHORIZATION, HeaderName, HeaderValue, InvalidHeaderName, InvalidHeaderValue}; -use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::{BTreeMap, HashMap}; -use std::num::NonZeroU16; -use std::path::Path; -use utils::matches_name_pattern; -use utils::variables::{VarError, Vars}; - -#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema)] -pub struct McpConfig { - #[serde(alias = "mcpServers")] - pub servers: BTreeMap, -} - -#[doc = include_str!("../docs/mcp_server_config.md")] -#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] -#[serde(untagged)] -pub enum McpServerConfig { - Stdio(StdioServerConfig), - Remote(RemoteServerConfig), - InMemory(InMemoryServerConfig), -} - -#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] -#[serde(deny_unknown_fields)] -pub struct StdioServerConfig { - /// Transport discriminant; always `stdio`. - #[serde(rename = "type", default)] - pub type_: StdioType, - - /// Executable launched to run the MCP server over stdio. - pub command: String, - - /// Command-line arguments passed to the executable. - #[serde(default)] - pub args: Vec, - - /// Environment variables set for the server process. - #[serde(default)] - pub env: HashMap, - - /// Controls which tools are deferred from the model-visible tool definitions. - #[serde(rename = "deferTools", alias = "proxy", default, skip_serializing_if = "ToolExposure::is_model_visible")] - pub defer_tools: ToolExposure, -} - -pub const AETHER_OAUTH_CLIENT_METADATA_URL: &str = "https://aether-agent.io/oauth/client-metadata.json"; -pub const AETHER_OAUTH_CALLBACK_PORT: NonZeroU16 = NonZeroU16::new(3118).unwrap(); - -#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct McpOAuthConfig { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub client_id: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub client_metadata_url: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub callback_port: Option, -} - -#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] -#[serde(deny_unknown_fields)] -pub struct RemoteServerConfig { - /// Transport discriminant; `http` (streamable HTTP) or `sse` (Server-Sent Events). - #[serde(rename = "type")] - pub type_: RemoteType, - - /// Base URL of the remote MCP server. - pub url: String, - - /// Extra HTTP headers sent with every request. - #[serde(default)] - pub headers: HashMap, - - /// OAuth settings for a pre-registered public client. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub oauth: Option, - - /// Controls which tools are deferred from the model-visible tool definitions. - #[serde(rename = "deferTools", alias = "proxy", default, skip_serializing_if = "ToolExposure::is_model_visible")] - pub defer_tools: ToolExposure, -} - -#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] -#[serde(deny_unknown_fields)] -pub struct InMemoryServerConfig { - /// Transport discriminant; always `in-memory`. - #[serde(rename = "type")] - pub type_: InMemoryType, - - /// Arguments passed to the built-in (in-process) server. - #[serde(default)] - pub args: Vec, - - /// Optional JSON input passed to the built-in server at startup. - #[serde(default)] - pub input: Option, - - /// Controls which tools are deferred from the model-visible tool definitions. - #[serde(rename = "deferTools", alias = "proxy", default, skip_serializing_if = "ToolExposure::is_model_visible")] - pub defer_tools: ToolExposure, -} - -#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, JsonSchema, PartialEq)] -pub enum StdioType { - #[default] - #[serde(rename = "stdio")] - Stdio, -} - -#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)] -pub enum RemoteType { - #[serde(rename = "http")] - Http, - #[serde(rename = "sse")] - Sse, -} - -#[derive(Debug, Clone, Copy, Deserialize, Serialize, JsonSchema, PartialEq)] -pub enum InMemoryType { - #[serde(rename = "in-memory")] - InMemory, -} - -/// Which of a server's tools are model-visible or deferred for progressive discovery. -#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize, JsonSchema)] -#[serde(from = "ToolExposureConfig", into = "ToolExposureConfig")] -#[schemars(with = "ToolExposureConfig")] -pub enum ToolExposure { - #[default] - ModelVisible, - Deferred(DeferredToolRules), -} - -#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] -#[serde(deny_unknown_fields)] -pub struct DeferredToolRules { - /// Tool names to defer. An empty list includes every tool. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub include: Vec, - - /// Tool names to keep model-visible. Exclude rules take precedence over include rules. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub exclude: Vec, -} - -#[derive(Debug, Clone)] -pub struct McpServer { - pub name: String, - pub transport: McpTransport, - pub tool_exposure: ToolExposure, -} - -#[derive(Debug, Clone)] -pub enum McpTransport { - Stdio { command: String, args: Vec, env: HashMap }, - Http(McpHttpConfig), - InMemory { spec: InMemoryServerSpec }, -} - -#[derive(Clone, Debug)] -pub struct InMemoryServerSpec { - pub factory: String, - pub args: Vec, - pub input: Option, -} - -#[derive(Debug, Clone)] -pub struct McpHttpConfig { - pub transport: StreamableHttpClientTransportConfig, - pub oauth: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ResolvedOAuth { - pub client_registration: OAuthClientRegistration, - pub callback_port: NonZeroU16, -} - -impl ResolvedOAuth { - pub fn redirect_uri(&self) -> String { - loopback_redirect_uri(self.callback_port.get()) - } -} - -pub fn loopback_redirect_uri(port: u16) -> String { - format!("http://localhost:{port}/") -} - -impl McpHttpConfig { - pub(crate) fn has_explicit_authorization(&self) -> bool { - self.transport.auth_header.is_some() || self.transport.custom_headers.contains_key(&AUTHORIZATION) - } - - pub fn resolved_oauth(&self) -> Option { - if self.has_explicit_authorization() { - return None; - } - let oauth = self.oauth.as_ref(); - let client_registration = match oauth { - Some(McpOAuthConfig { client_id: Some(client_id), .. }) => { - OAuthClientRegistration::PreRegistered(client_id.clone()) - } - Some(McpOAuthConfig { client_metadata_url: Some(url), .. }) => { - OAuthClientRegistration::ClientMetadata(url.clone()) - } - _ => OAuthClientRegistration::ClientMetadata(AETHER_OAUTH_CLIENT_METADATA_URL.to_string()), - }; - Some(ResolvedOAuth { - client_registration, - callback_port: oauth.and_then(|oauth| oauth.callback_port).unwrap_or(AETHER_OAUTH_CALLBACK_PORT), - }) - } -} - -impl From for McpHttpConfig { - fn from(transport: StreamableHttpClientTransportConfig) -> Self { - Self { transport, oauth: None } - } -} - -impl ToolExposure { - pub fn deferred_all() -> Self { - Self::Deferred(DeferredToolRules::default()) - } - - pub fn is_model_visible(&self) -> bool { - matches!(self, Self::ModelVisible) - } - - pub fn has_deferred_tools(&self) -> bool { - matches!(self, Self::Deferred(_)) - } - - pub fn is_model_visible_tool(&self, tool_name: &str) -> bool { - match self { - Self::ModelVisible => true, - Self::Deferred(rules) => !rules.matches(tool_name), - } - } - - /// Defer every tool, preserving any existing rules. - pub fn defer_all_tools(&mut self) { - if self.is_model_visible() { - *self = Self::deferred_all(); - } - } -} - -impl DeferredToolRules { - pub fn new(include: &[&str], exclude: &[&str]) -> Self { - Self { - include: include.iter().map(ToString::to_string).collect(), - exclude: exclude.iter().map(ToString::to_string).collect(), - } - } - - fn matches(&self, tool_name: &str) -> bool { - let included = - self.include.is_empty() || self.include.iter().any(|pattern| matches_name_pattern(pattern, tool_name)); - let excluded = self.exclude.iter().any(|pattern| matches_name_pattern(pattern, tool_name)); - included && !excluded - } -} - -/// The `deferTools` config field's wire shape: a boolean or an include/exclude object. -#[derive(Deserialize, Serialize, JsonSchema)] -#[serde(untagged)] -enum ToolExposureConfig { - Enabled(bool), - Rules(DeferredToolRules), -} - -impl From for ToolExposure { - fn from(repr: ToolExposureConfig) -> Self { - match repr { - ToolExposureConfig::Enabled(false) => Self::ModelVisible, - ToolExposureConfig::Enabled(true) => Self::deferred_all(), - ToolExposureConfig::Rules(rules) => Self::Deferred(rules), - } - } -} - -impl From for ToolExposureConfig { - fn from(exposure: ToolExposure) -> Self { - match exposure { - ToolExposure::ModelVisible => Self::Enabled(false), - ToolExposure::Deferred(rules) if rules == DeferredToolRules::default() => Self::Enabled(true), - ToolExposure::Deferred(rules) => Self::Rules(rules), - } - } -} - -impl McpServer { - pub fn new(name: impl Into, transport: McpTransport, tool_exposure: ToolExposure) -> Self { - Self { name: name.into(), transport, tool_exposure } - } - - pub fn with_exposure(mut self, exposure: ToolExposure) -> Self { - self.tool_exposure = exposure; - self - } - - pub fn has_deferred_tools(&self) -> bool { - self.tool_exposure.has_deferred_tools() - } -} - -#[derive(Debug, thiserror::Error)] -pub enum ParseError { - #[error("Failed to read config file: {0}")] - IoError(#[from] std::io::Error), - - #[error("Invalid JSON: {0}")] - JsonError(#[from] serde_json::Error), - - #[error("Variable expansion failed: {0}")] - VarError(#[from] VarError), - - #[error("Invalid HTTP header name {name:?}: {source}")] - InvalidHeaderName { - name: String, - #[source] - source: InvalidHeaderName, - }, - - #[error("Invalid HTTP header value for {name:?}: {source}")] - InvalidHeaderValue { - name: String, - #[source] - source: InvalidHeaderValue, - }, -} - -impl McpConfig { - pub fn new(servers: BTreeMap) -> Self { - Self { servers } - } - - pub fn from_json_file(path: impl AsRef) -> Result { - let content = std::fs::read_to_string(path)?; - Self::from_json(&content) - } - - pub fn from_json_files>(paths: &[T]) -> Result { - let mut merged = BTreeMap::new(); - for path in paths { - let raw = Self::from_json_file(path)?; - merged.extend(raw.servers); - } - Ok(Self::new(merged)) - } - - pub fn from_json(json: &str) -> Result { - Ok(serde_json::from_str(json)?) - } - - pub fn into_servers(self, vars: &Vars) -> Result, ParseError> { - self.into_servers_with_deferred_tools(vars, false) - } - - pub fn into_servers_with_deferred_tools( - self, - vars: &Vars, - defer_all_tools: bool, - ) -> Result, ParseError> { - self.servers.into_iter().map(|(name, config)| config.into_server(name, vars, defer_all_tools)).collect() - } - - pub fn defer_all_tools(&mut self) { - for server in self.servers.values_mut() { - server.defer_all_tools(); - } - } -} - -impl McpServerConfig { - pub fn defer_tools(&self) -> &ToolExposure { - match self { - McpServerConfig::Stdio(config) => &config.defer_tools, - McpServerConfig::Remote(config) => &config.defer_tools, - McpServerConfig::InMemory(config) => &config.defer_tools, - } - } - - pub fn defer_all_tools(&mut self) { - let exposure = match self { - McpServerConfig::Stdio(config) => &mut config.defer_tools, - McpServerConfig::Remote(config) => &mut config.defer_tools, - McpServerConfig::InMemory(config) => &mut config.defer_tools, - }; - exposure.defer_all_tools(); - } - - pub fn into_server(self, name: String, vars: &Vars, defer_all_tools: bool) -> Result { - let mut exposure = self.defer_tools().clone(); - if defer_all_tools { - exposure.defer_all_tools(); - } - let transport = self.into_transport(name.clone(), vars)?; - Ok(McpServer { name, transport, tool_exposure: exposure }) - } - - fn into_transport(self, name: String, vars: &Vars) -> Result { - match self { - McpServerConfig::Stdio(StdioServerConfig { command, args, env, .. }) => Ok(McpTransport::Stdio { - command: vars.expand(&command)?, - args: args.into_iter().map(|a| vars.expand(&a)).collect::, _>>()?, - env: env - .into_iter() - .map(|(k, v)| Ok((k, vars.expand(&v)?))) - .collect::, VarError>>()?, - }), - - McpServerConfig::Remote(RemoteServerConfig { url, headers, oauth, .. }) => { - let mut transport = StreamableHttpClientTransportConfig::with_uri(vars.expand(&url)?); - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|source| ParseError::InvalidHeaderName { name: name.clone(), source })?; - - let mut header_value = vars - .expand(&value)? - .parse::() - .map_err(|source| ParseError::InvalidHeaderValue { name, source })?; - - if header_name == AUTHORIZATION { - header_value.set_sensitive(true); - } - - transport.custom_headers.insert(header_name, header_value); - } - - let oauth = oauth - .map(|oauth| -> Result { - Ok(McpOAuthConfig { - client_id: oauth.client_id.map(|value| vars.expand(&value)).transpose()?, - client_metadata_url: oauth - .client_metadata_url - .map(|value| vars.expand(&value)) - .transpose()?, - callback_port: oauth.callback_port, - }) - }) - .transpose()?; - - Ok(McpTransport::Http(McpHttpConfig { transport, oauth })) - } - - McpServerConfig::InMemory(InMemoryServerConfig { args, input, .. }) => { - let args = args.into_iter().map(|a| vars.expand(&a)).collect::, VarError>>()?; - Ok(McpTransport::InMemory { spec: InMemoryServerSpec { factory: name, args, input } }) - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - use tempfile::tempdir; - - fn write_config(dir: &Path, name: &str, json: &str) -> std::path::PathBuf { - let path = dir.join(name); - fs::write(&path, json).unwrap(); - path - } - - fn stdio_config(command: &str) -> String { - format!(r#"{{"servers": {{"coding": {{"type": "stdio", "command": "{command}"}}}}}}"#) - } - - #[test] - fn from_json_accepts_mcp_servers_key() { - let config = McpConfig::from_json(r#"{"mcpServers": {"alpha": {"type": "stdio", "command": "a"}}}"#).unwrap(); - assert_eq!(config.servers.len(), 1); - assert!(config.servers.contains_key("alpha")); - } - - #[test] - fn from_json_defaults_missing_type_to_stdio() { - let config = McpConfig::from_json( - r#"{"mcpServers": {"devtools": {"command": "npx", "args": ["-y", "chrome-devtools-mcp"]}}}"#, - ) - .unwrap(); - match config.servers.get("devtools").unwrap() { - McpServerConfig::Stdio(StdioServerConfig { command, args, defer_tools: exposure, .. }) => { - assert_eq!(command, "npx"); - assert_eq!(args, &["-y", "chrome-devtools-mcp"]); - assert!(exposure.is_model_visible()); - } - other => panic!("expected Stdio server, got {other:?}"), - } - } - - #[test] - fn from_json_accepts_legacy_server_proxy_true() { - let config = - McpConfig::from_json(r#"{"servers": {"playwright": {"type": "stdio", "command": "npx", "proxy": true}}}"#) - .unwrap(); - assert!(config.servers.get("playwright").unwrap().defer_tools().has_deferred_tools()); - } - - #[test] - fn from_json_accepts_server_defer_tools_true() { - let config = McpConfig::from_json( - r#"{"servers": {"playwright": {"type": "stdio", "command": "npx", "deferTools": true}}}"#, - ) - .unwrap(); - assert!(config.servers.get("playwright").unwrap().defer_tools().has_deferred_tools()); - } - - #[test] - fn from_json_rejects_unknown_server_type() { - let result = McpConfig::from_json(r#"{"servers":{"tools":{"type":"deferTools","servers":{}}}}"#); - assert!(result.is_err()); - } - - #[test] - fn false_defer_tools_omits_during_serialization() { - let config = - McpConfig::from_json(r#"{"servers": {"coding": {"type": "stdio", "command": "a", "deferTools": false}}}"#) - .unwrap(); - let serialized = serde_json::to_string(&config).unwrap(); - assert!(!serialized.contains("deferTools")); - } - - #[test] - fn true_defer_tools_serializes() { - let config = - McpConfig::from_json(r#"{"servers": {"coding": {"type": "stdio", "command": "a", "deferTools": true}}}"#) - .unwrap(); - let serialized = serde_json::to_string(&config).unwrap(); - assert!(serialized.contains("deferTools")); - } - - #[test] - fn from_json_rejects_unknown_type() { - let result = McpConfig::from_json(r#"{"servers": {"bad": {"type": "htp", "url": "https://example.com"}}}"#); - assert!(result.is_err()); - } - - #[test] - fn from_json_files_empty_returns_empty_servers() { - let result = McpConfig::from_json_files::<&str>(&[]).unwrap(); - assert!(result.servers.is_empty()); - } - - #[test] - fn from_json_files_single_file_matches_from_json_file() { - let dir = tempdir().unwrap(); - let path = write_config(dir.path(), "a.json", &stdio_config("ls")); - - let single = McpConfig::from_json_file(&path).unwrap(); - let multi = McpConfig::from_json_files(&[&path]).unwrap(); - - assert_eq!(single.servers.len(), multi.servers.len()); - assert!(multi.servers.contains_key("coding")); - } - - #[test] - fn from_json_files_merges_disjoint_servers() { - let dir = tempdir().unwrap(); - let a = write_config(dir.path(), "a.json", r#"{"servers": {"alpha": {"type": "stdio", "command": "a"}}}"#); - let b = write_config(dir.path(), "b.json", r#"{"servers": {"beta": {"type": "stdio", "command": "b"}}}"#); - - let merged = McpConfig::from_json_files(&[a, b]).unwrap(); - assert_eq!(merged.servers.len(), 2); - assert!(merged.servers.contains_key("alpha")); - assert!(merged.servers.contains_key("beta")); - } - - #[test] - fn from_json_rejects_unknown_exposure_fields_for_all_transports() { - for server in [ - r#"{"command":"x","direct_tool":["bash"]}"#, - r#"{"type":"http","url":"https://example.com","direct_tool":["bash"]}"#, - r#"{"type":"in-memory","direct_tool":["bash"]}"#, - ] { - let json = format!(r#"{{"servers":{{"bad":{server}}}}}"#); - assert!(McpConfig::from_json(&json).is_err(), "unknown field was accepted: {server}"); - } - } - - #[test] - fn from_json_files_last_file_wins_on_collision_including_exposure() { - let dir = tempdir().unwrap(); - let a = write_config( - dir.path(), - "a.json", - r#"{"servers":{"coding":{"type":"stdio","command":"from_a","deferTools":{"exclude":["bash"]}}}}"#, - ); - let b = write_config(dir.path(), "b.json", r#"{"servers":{"coding":{"type":"stdio","command":"from_b"}}}"#); - - let merged_ab = McpConfig::from_json_files(&[&a, &b]).unwrap(); - match merged_ab.servers.get("coding").unwrap() { - McpServerConfig::Stdio(StdioServerConfig { command, defer_tools: exposure, .. }) => { - assert_eq!(command, "from_b"); - assert_eq!(exposure, &ToolExposure::ModelVisible); - } - other => panic!("expected Stdio, got {other:?}"), - } - - let merged_ba = McpConfig::from_json_files(&[&b, &a]).unwrap(); - match merged_ba.servers.get("coding").unwrap() { - McpServerConfig::Stdio(StdioServerConfig { command, defer_tools: exposure, .. }) => { - assert_eq!(command, "from_a"); - assert_eq!(exposure, &ToolExposure::Deferred(DeferredToolRules::new(&[], &["bash"]))); - } - other => panic!("expected Stdio, got {other:?}"), - } - } - - #[test] - fn defer_all_tools_sets_every_server() { - let mut config = McpConfig::from_json( - r#"{"servers":{"a":{"type":"stdio","command":"a"},"b":{"type":"http","url":"https://example.com"}}}"#, - ) - .unwrap(); - config.defer_all_tools(); - assert!(config.servers.values().all(|server| server.defer_tools().has_deferred_tools())); - } - - #[test] - fn from_json_files_propagates_io_error_on_missing_file() { - let dir = tempdir().unwrap(); - let missing = dir.path().join("does-not-exist.json"); - let result = McpConfig::from_json_files(&[missing]); - assert!(matches!(result, Err(ParseError::IoError(_)))); - } - - #[test] - fn from_json_files_propagates_json_error_on_invalid_file() { - let dir = tempdir().unwrap(); - let bad = write_config(dir.path(), "bad.json", "not valid json"); - let result = McpConfig::from_json_files(&[bad]); - assert!(matches!(result, Err(ParseError::JsonError(_)))); - } - - #[tokio::test] - async fn into_servers_preserves_deferred_tool_flags() { - let json = r#"{ - "servers": { - "github": {"type": "stdio", "command": "g"}, - "playwright": {"type": "stdio", "command": "p", "deferTools": true} - } - }"#; - let config = McpConfig::from_json(json).unwrap(); - let servers = config.into_servers(&Vars::new()).unwrap(); - - assert_eq!(servers.len(), 2); - assert!(!servers.iter().find(|s| s.name == "github").unwrap().has_deferred_tools()); - assert!(servers.iter().find(|s| s.name == "playwright").unwrap().has_deferred_tools()); - } - - #[tokio::test] - async fn into_servers_with_deferred_tools_forces_deferred_tool_flags() { - let config = - McpConfig::from_json(r#"{"servers":{"github":{"type":"stdio","command":"g","deferTools":false}}}"#) - .unwrap(); - let servers = config.into_servers_with_deferred_tools(&Vars::new(), true).unwrap(); - assert!(servers[0].has_deferred_tools()); - } - - #[test] - fn defer_tools_accepts_boolean_or_rules_for_all_transport_shapes() { - let config = McpConfig::from_json( - r#"{"servers":{"all":{"command":"a","deferTools":true},"stdio":{"command":"x","deferTools":{"include":["lsp_*"],"exclude":["lsp_rename"]}},"http":{"type":"http","url":"https://example.com","deferTools":{"exclude":["bash"]}},"memory":{"type":"in-memory","deferTools":{"include":["read"]}}}}"#, - ) - .unwrap(); - - assert_eq!(config.servers["all"].defer_tools(), &ToolExposure::deferred_all()); - assert_eq!( - config.servers["stdio"].defer_tools(), - &ToolExposure::Deferred(DeferredToolRules::new(&["lsp_*"], &["lsp_rename"])) - ); - assert_eq!( - config.servers["http"].defer_tools(), - &ToolExposure::Deferred(DeferredToolRules::new(&[], &["bash"])) - ); - assert_eq!( - config.servers["memory"].defer_tools(), - &ToolExposure::Deferred(DeferredToolRules::new(&["read"], &[])) - ); - } - - #[test] - fn deferred_tool_rules_serialize_and_defaults_are_omitted() { - let config = McpConfig::from_json( - r#"{"servers":{"coding":{"command":"x","deferTools":{"exclude":["bash","lsp_*"]}},"direct":{"command":"y"},"full":{"command":"z","deferTools":true}}}"#, - ) - .unwrap(); - let value = serde_json::to_value(config).unwrap(); - - assert_eq!(value["servers"]["coding"]["deferTools"], serde_json::json!({"exclude":["bash", "lsp_*"]})); - assert!(value["servers"]["direct"].get("deferTools").is_none()); - assert_eq!(value["servers"]["full"]["deferTools"], serde_json::json!(true)); - } - - #[test] - fn legacy_direct_tools_is_rejected() { - let result = - McpConfig::from_json(r#"{"servers":{"coding":{"command":"x","deferTools":true,"direct_tools":["bash"]}}}"#); - assert!(result.is_err()); - } - - #[tokio::test] - async fn deferred_tool_rules_partition_tools_with_exclude_winning() { - let config = McpConfig::from_json( - r#"{"servers":{"coding":{"command":"server","deferTools":{"include":["lsp_*","bash"],"exclude":["lsp_rename"]}}}}"#, - ) - .unwrap(); - let servers = config.into_servers(&Vars::new()).unwrap(); - let exposure = &servers[0].tool_exposure; - - assert!(!exposure.is_model_visible_tool("lsp_hover")); - assert!(exposure.is_model_visible_tool("lsp_rename")); - assert!(!exposure.is_model_visible_tool("bash")); - assert!(exposure.is_model_visible_tool("read_file")); - } - - #[tokio::test] - async fn forced_deferral_preserves_per_server_rules() { - let config = - McpConfig::from_json(r#"{"servers":{"coding":{"command":"server","deferTools":{"exclude":["bash"]}}}}"#) - .unwrap(); - let servers = config.into_servers_with_deferred_tools(&Vars::new(), true).unwrap(); - assert!(servers[0].has_deferred_tools()); - assert!(servers[0].tool_exposure.is_model_visible_tool("bash")); - assert!(!servers[0].tool_exposure.is_model_visible_tool("read_file")); - } - - #[tokio::test] - async fn into_transport_expands_workspace_var_in_stdio_args() { - let config = McpConfig::from_json( - r#"{"servers":{"coding":{"type":"stdio","command":"server","args":["--root","${WORKSPACE}/src"]}}}"#, - ) - .unwrap(); - let vars = Vars::new().with("WORKSPACE", "/workspace"); - let servers = config.into_servers(&vars).unwrap(); - - match &servers[0].transport { - McpTransport::Stdio { args, .. } => { - assert_eq!(args, &["--root", "/workspace/src"]); - } - other => panic!("expected Stdio transport, got {other:?}"), - } - } - - #[tokio::test] - async fn into_transport_preserves_bearer_auth_header() -> Result<(), String> { - let config = McpConfig::from_json( - r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"Authorization":"Bearer secret-token"}}}}"#, - ) - .map_err(|e| e.to_string())?; - - let servers = config.into_servers(&Vars::new()).map_err(|e| e.to_string())?; - let McpTransport::Http(config) = &servers[0].transport else { - return Err(format!("expected Http transport, got {:?}", servers[0].transport)); - }; - - assert!(config.transport.auth_header.is_none()); - assert_eq!(config.transport.custom_headers[&AUTHORIZATION], "Bearer secret-token"); - Ok(()) - } - - #[tokio::test] - async fn into_transport_keeps_non_bearer_auth_header_verbatim() -> Result<(), String> { - let config = McpConfig::from_json( - r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"Authorization":"Basic dXNlcjpwYXNz"}}}}"#, - ) - .map_err(|e| e.to_string())?; - let servers = config.into_servers(&Vars::new()).map_err(|e| e.to_string())?; - - let McpTransport::Http(config) = &servers[0].transport else { - return Err(format!("expected Http transport, got {:?}", servers[0].transport)); - }; - assert!(config.transport.auth_header.is_none()); - assert_eq!(config.transport.custom_headers[&AUTHORIZATION], "Basic dXNlcjpwYXNz"); - Ok(()) - } - - #[tokio::test] - async fn into_transport_expands_vars_in_auth_header() -> Result<(), String> { - let config = McpConfig::from_json( - r#"{"servers":{"weather":{"type":"http","url":"http://127.0.0.1:9000/mcp","headers":{"Authorization":"Bearer ${TOKEN}"}}}}"#, - ) - .map_err(|e| e.to_string())?; - let vars = Vars::new().with("TOKEN", "expanded-token"); - let servers = config.into_servers(&vars).map_err(|e| e.to_string())?; - - let McpTransport::Http(config) = &servers[0].transport else { - return Err(format!("expected Http transport, got {:?}", servers[0].transport)); - }; - assert!(config.transport.auth_header.is_none()); - assert_eq!(config.transport.custom_headers[&AUTHORIZATION], "Bearer expanded-token"); - Ok(()) - } -} diff --git a/crates/mcp-utils/src/client/connection.rs b/crates/mcp-utils/src/client/connection.rs deleted file mode 100644 index 5e7f821c4..000000000 --- a/crates/mcp-utils/src/client/connection.rs +++ /dev/null @@ -1,385 +0,0 @@ -use super::{ - McpClientEvent, McpError, OAuthHandlerFactory, Result, - config::McpHttpConfig, - manager::{RuntimeMcpServer, RuntimeMcpTransport, ToolListChangedRequest}, - mcp_client::McpClient, -}; -use crate::{client::OAuthHandlerContext, protocol::client_lifecycle_mode}; -use aether_auth::{OAuthCredentialStorage, create_auth_manager_from_store, perform_oauth_flow}; -use llm::ToolAnnotations; -use rmcp::{ - RoleClient, RoleServer, ServiceExt, - model::{ClientConfig, Tool as RmcpTool}, - serve_client_with_lifecycle, - service::{DynService, RunningService}, - transport::{ - StreamableHttpClientTransport, TokioChildProcess, auth::AuthClient, - streamable_http_client::StreamableHttpClientTransportConfig, - }, -}; -use serde_json::Value; -use std::collections::HashMap; -use std::path::PathBuf; -use std::process::Stdio; -use std::sync::{ - Arc, - atomic::{AtomicU64, Ordering}, -}; -use tokio::{ - io::{AsyncBufReadExt, BufReader}, - process::{ChildStderr, Command}, - sync::mpsc, - task::JoinHandle, -}; - -#[derive(Debug, Clone)] -pub struct Tool { - pub name: String, - pub description: String, - pub parameters: Value, - pub annotations: Option, -} - -pub(crate) fn convert_tool_annotations(annotations: &rmcp::model::ToolAnnotations) -> ToolAnnotations { - ToolAnnotations { - title: annotations.title.clone(), - read_only_hint: annotations.read_only_hint, - destructive_hint: annotations.destructive_hint, - idempotent_hint: annotations.idempotent_hint, - open_world_hint: annotations.open_world_hint, - } -} - -impl From for Tool { - fn from(tool: RmcpTool) -> Self { - Self::from(&tool) - } -} - -impl From<&RmcpTool> for Tool { - fn from(tool: &RmcpTool) -> Self { - Self { - name: tool.name.to_string(), - description: tool.description.clone().unwrap_or_default().to_string(), - parameters: serde_json::Value::Object((*tool.input_schema).clone()), - annotations: tool.annotations.as_ref().map(convert_tool_annotations), - } - } -} - -pub(super) struct ConnectConfig { - pub client_info: ClientConfig, - pub event_sender: mpsc::Sender, - pub tool_refresh_sender: mpsc::Sender, - pub next_connection_generation: Arc, - pub root_dir: PathBuf, - pub oauth_handler_factory: Option, - pub oauth_credential_store: Option>, -} - -/// The result of attempting to connect (or authenticate) to an MCP server. -pub struct McpConnectAttempt { - pub name: String, - pub outcome: McpConnectOutcome, -} - -pub enum McpConnectOutcome { - Connected { conn: McpServerConnection, reauth_config: Option }, - NeedsOAuth { config: McpHttpConfig, challenge: Option, error: McpError }, - Failed { error: McpError }, -} - -impl McpConnectAttempt { - pub fn failed(name: impl Into, error: McpError) -> Self { - Self { name: name.into(), outcome: McpConnectOutcome::Failed { error } } - } -} - -pub struct McpServerConnection { - pub(super) client: Arc>, - pub(super) server_task: Option>, - pub(super) instructions: Option, - generation: u64, -} - -impl McpServerConnection { - pub(super) async fn reconnect_with_auth( - name: &str, - config: StreamableHttpClientTransportConfig, - auth_client: AuthClient, - mcp_client: McpClient, - generation: u64, - ) -> Result { - let transport = StreamableHttpClientTransport::with_client(auth_client, config); - let client = serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()) - .await - .map_err(|e| McpError::ConnectionFailed(format!("reconnect failed for '{name}': {e}")))?; - Ok(Self::from_parts(client, None, generation)) - } - - pub(super) async fn list_tools(&self) -> Result> { - self.client - .list_all_tools() - .await - .map_err(|e| McpError::ToolDiscoveryFailed(format!("Failed to list tools: {e}"))) - } - - fn from_parts( - client: RunningService, - server_task: Option>, - generation: u64, - ) -> Self { - let instructions = client.peer_info().and_then(|info| info.instructions.clone()).filter(|s| !s.is_empty()); - Self { client: Arc::new(client), server_task, instructions, generation } - } - - pub(super) fn generation(&self) -> u64 { - self.generation - } -} - -pub(super) async fn connect_server(server: RuntimeMcpServer, ctx: &ConnectConfig) -> McpConnectAttempt { - let RuntimeMcpServer { name, transport, tool_exposure: _ } = server; - let reauth_config = reauth_config_for(&transport, ctx.oauth_handler_factory.as_ref()); - let generation = ctx.next_connection_generation.fetch_add(1, Ordering::Relaxed); - let mcp_client = McpClient::new(ctx.client_info.clone(), name.clone(), ctx.event_sender.clone()) - .with_tool_refresh(ctx.tool_refresh_sender.clone(), generation); - - let outcome = match transport { - RuntimeMcpTransport::Stdio { command, args, env } => { - connect_stdio(&name, command, args, env, mcp_client, ctx.root_dir.clone(), generation).await - } - RuntimeMcpTransport::InMemory { server } => connect_in_memory(&name, server, mcp_client, generation).await, - RuntimeMcpTransport::Http(config) => { - connect_http( - &name, - config, - mcp_client, - ctx.oauth_handler_factory.as_ref(), - ctx.oauth_credential_store.as_ref(), - generation, - ) - .await - } - }; - - McpConnectAttempt { name, outcome: outcome.with_reauth(reauth_config) } -} - -pub async fn authenticate_http( - name: String, - config: McpHttpConfig, - challenge: Option, - ctx: Arc, -) -> McpConnectAttempt { - let outcome = match async { - let factory = ctx - .oauth_handler_factory - .as_ref() - .ok_or_else(|| McpError::ConnectionFailed(format!("No OAuth handler factory available for '{name}'")))?; - let oauth = config - .resolved_oauth() - .ok_or_else(|| McpError::ConnectionFailed(format!("OAuth is not available for '{name}'")))?; - let handler = factory(OAuthHandlerContext { - server_name: name.clone(), - callback_port: Some(oauth.callback_port), - tx: ctx.event_sender.clone(), - })?; - - let auth_client = perform_oauth_flow( - &name, - &config.transport.uri, - handler.as_ref(), - aether_auth::OAuthFlowOptions { client_registration: oauth.client_registration, challenge }, - ctx.oauth_credential_store.clone(), - ) - .await - .map_err(|e| McpError::ConnectionFailed(format!("OAuth failed for '{name}': {e}")))?; - - let generation = ctx.next_connection_generation.fetch_add(1, Ordering::Relaxed); - let mcp_client = McpClient::new(ctx.client_info.clone(), name.clone(), ctx.event_sender.clone()) - .with_tool_refresh(ctx.tool_refresh_sender.clone(), generation); - McpServerConnection::reconnect_with_auth(&name, config.transport.clone(), auth_client, mcp_client, generation) - .await - } - .await - { - Ok(conn) => McpConnectOutcome::Connected { conn, reauth_config: Some(config) }, - Err(error) => McpConnectOutcome::Failed { error }, - }; - - McpConnectAttempt { name, outcome } -} - -impl McpConnectOutcome { - fn with_reauth(self, reauth_config: Option) -> Self { - match self { - Self::Connected { conn, .. } => Self::Connected { conn, reauth_config }, - other => other, - } - } -} - -async fn connect_stdio( - server_name: &str, - command: String, - args: Vec, - env: HashMap, - mcp_client: McpClient, - cwd: PathBuf, - generation: u64, -) -> McpConnectOutcome { - let mut cmd = Command::new(&command); - cmd.args(&args).envs(&env).current_dir(&cwd); - - let (proc, stderr) = match TokioChildProcess::builder(cmd).stderr(Stdio::piped()).spawn() { - Ok(parts) => parts, - Err(e) => return McpConnectOutcome::Failed { error: McpError::SpawnFailed { command, reason: e.to_string() } }, - }; - let stderr_task = stderr.map(|stderr| spawn_stderr_logger(server_name.to_string(), stderr)); - - match serve_client_with_lifecycle(mcp_client, proc, client_lifecycle_mode()).await { - Ok(client) => McpConnectOutcome::Connected { - conn: McpServerConnection::from_parts(client, stderr_task, generation), - reauth_config: None, - }, - Err(e) => { - if let Some(task) = stderr_task { - task.abort(); - } - McpConnectOutcome::Failed { error: McpError::from(e) } - } - } -} - -fn spawn_stderr_logger(server_name: String, stderr: ChildStderr) -> JoinHandle<()> { - tokio::spawn(async move { - let mut lines = BufReader::new(stderr).lines(); - loop { - match lines.next_line().await { - Ok(Some(line)) => tracing::info!(server = %server_name, stderr = %line, "MCP server stderr"), - Ok(None) => break, - Err(error) => { - tracing::warn!(server = %server_name, %error, "failed to read MCP server stderr"); - break; - } - } - } - }) -} - -async fn connect_in_memory( - name: &str, - server: Box>, - mcp_client: McpClient, - generation: u64, -) -> McpConnectOutcome { - match serve_in_memory(server, mcp_client, name).await { - Ok((client, handle)) => McpConnectOutcome::Connected { - conn: McpServerConnection::from_parts(client, Some(handle), generation), - reauth_config: None, - }, - Err(error) => McpConnectOutcome::Failed { error }, - } -} - -async fn connect_http( - name: &str, - config: McpHttpConfig, - mcp_client: McpClient, - oauth_handler_factory: Option<&OAuthHandlerFactory>, - oauth_credential_store: Option<&Arc>, - generation: u64, -) -> McpConnectOutcome { - let conn_err = |e| McpError::ConnectionFailed(format!("HTTP MCP server {name}: {e}")); - let oauth = config.resolved_oauth(); - let restored = if let (Some(store), Some(oauth)) = (oauth_credential_store, oauth.as_ref()) { - match create_auth_manager_from_store( - name, - &config.transport.uri, - oauth.client_registration.pre_registered_client_id(), - &oauth.redirect_uri(), - Arc::clone(store), - ) - .await - { - Ok(manager) => manager, - Err(e) => { - tracing::warn!( - server = %name, - error = %e, - "Failed to initialize auth manager from stored credentials, proceeding without auth" - ); - None - } - } - } else { - None - }; - let result = if let Some(manager) = restored { - tracing::debug!("Using OAuth for server '{name}'"); - let auth_client = AuthClient::new(reqwest::Client::default(), manager); - let transport = StreamableHttpClientTransport::with_client(auth_client, config.transport.clone()); - serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()).await - } else { - let transport = StreamableHttpClientTransport::from_config(config.transport.clone()); - serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()).await - }; - - match result { - Ok(client) => McpConnectOutcome::Connected { - conn: McpServerConnection::from_parts(client, None, generation), - reauth_config: None, - }, - Err(error) => { - let challenge = error.auth_challenge().map(str::to_string); - let authorization_required = error.is_authorization_required(); - let error = conn_err(error); - tracing::warn!("Failed to connect to MCP server '{name}': {error}"); - if oauth_handler_factory.is_some() && oauth.is_some() && (authorization_required || challenge.is_some()) { - McpConnectOutcome::NeedsOAuth { config, challenge, error } - } else { - McpConnectOutcome::Failed { error } - } - } - } -} - -fn reauth_config_for( - transport: &RuntimeMcpTransport, - oauth_handler_factory: Option<&OAuthHandlerFactory>, -) -> Option { - match transport { - RuntimeMcpTransport::Http(config) - if oauth_handler_factory.is_some() && !config.has_explicit_authorization() => - { - Some(config.clone()) - } - _ => None, - } -} - -async fn serve_in_memory( - server: Box>, - mcp_client: McpClient, - label: &str, -) -> Result<(RunningService, JoinHandle<()>)> { - let (client_transport, server_transport) = tokio::io::duplex(64 * 1024); - - let server_handle = tokio::spawn(async move { - match server.serve(server_transport).await { - Ok(_service) => { - std::future::pending::<()>().await; - } - Err(e) => { - eprintln!("MCP server error: {e}"); - } - } - }); - - let client = serve_client_with_lifecycle(mcp_client, client_transport, client_lifecycle_mode()) - .await - .map_err(|e| McpError::ConnectionFailed(format!("Failed to connect to in-memory server '{label}': {e}")))?; - - Ok((client, server_handle)) -} diff --git a/crates/mcp-utils/src/client/connection_attempt_manager.rs b/crates/mcp-utils/src/client/connection_attempt_manager.rs deleted file mode 100644 index 775a45dbc..000000000 --- a/crates/mcp-utils/src/client/connection_attempt_manager.rs +++ /dev/null @@ -1,39 +0,0 @@ -use super::connection::McpConnectAttempt; -use std::collections::HashMap; -use tokio::task::{AbortHandle, JoinError, JoinSet}; - -#[derive(Default)] -pub struct McpConnectionAttemptManager { - set: JoinSet, - by_server: HashMap, -} - -impl McpConnectionAttemptManager { - /// Spawn a fresh auth task for `server`, aborting any in-flight task for - /// the same server so the user can retry a stuck OAuth flow. - pub fn spawn(&mut self, server: String, fut: impl Future + Send + 'static) { - if let Some(prior) = self.by_server.remove(&server) { - prior.abort(); - } - let handle = self.set.spawn(fut); - self.by_server.insert(server, handle); - } - - pub fn is_empty(&self) -> bool { - self.set.is_empty() - } - - pub async fn join_next(&mut self) -> Option> { - let joined = self.set.join_next().await; - if let Some(Ok(attempt)) = &joined { - self.by_server.remove(&attempt.name); - } - joined - } - - pub async fn shutdown(&mut self) { - self.set.abort_all(); - while self.set.join_next().await.is_some() {} - self.by_server.clear(); - } -} diff --git a/crates/mcp-utils/src/client/elicitation.rs b/crates/mcp-utils/src/client/elicitation.rs index d4f8372e5..cbc4d1814 100644 --- a/crates/mcp-utils/src/client/elicitation.rs +++ b/crates/mcp-utils/src/client/elicitation.rs @@ -1,40 +1,52 @@ -use crate::client::McpClient; -use rmcp::model::{ElicitRequestParams, ElicitResult, InputRequest, InputRequests, InputResponses, RequestMetaObject}; +use rmcp::model::{ + ElicitRequestParams, ElicitResult, ElicitationAction, ElicitationCapability, FormElicitationCapability, + UrlElicitationCapability, +}; +use tokio::sync::{mpsc, oneshot}; -pub(crate) enum ElicitInputsError { - UnsupportedInput, - Serialize(serde_json::Error), +#[derive(Debug, Clone)] +pub struct ElicitationSink { + pub(super) events: mpsc::Sender, + pub(super) capability: ElicitationCapability, } -pub(crate) async fn elicit_inputs( - client: &McpClient, - requests: InputRequests, -) -> Result<(InputResponses, Vec), ElicitInputsError> { - let mut responses = InputResponses::new(); - let mut results = Vec::new(); - for (key, request) in requests { - let InputRequest::Elicitation(elicitation_request) = request else { - return Err(ElicitInputsError::UnsupportedInput); - }; - let extension_meta = elicitation_request.extensions.get::().cloned(); - let params = with_meta(elicitation_request.params, extension_meta); - let result = client.dispatch_elicitation(params).await; - let response = serde_json::to_value(&result).map_err(ElicitInputsError::Serialize)?; - responses.insert(key, response); - results.push(result); +#[derive(Debug)] +pub enum Elicitation { + Request(Box), + Complete { server: String, id: String }, +} + +#[derive(Debug)] +pub struct ElicitationRequest { + pub server: String, + pub request: ElicitRequestParams, + responder: oneshot::Sender, +} + +impl ElicitationSink { + pub fn new(events: mpsc::Sender, capability: ElicitationCapability) -> Self { + Self { events, capability } + } + + pub fn form_and_url(events: mpsc::Sender) -> Self { + let capability = ElicitationCapability::new() + .with_form(FormElicitationCapability::new()) + .with_url(UrlElicitationCapability::new()); + Self::new(events, capability) } - Ok((responses, results)) } -pub(crate) fn with_meta(mut request: ElicitRequestParams, meta: Option) -> ElicitRequestParams { - if let Some(meta) = meta { - match &mut request { - ElicitRequestParams::FormElicitationParams { meta: request_meta, .. } - | ElicitRequestParams::UrlElicitationParams { meta: request_meta, .. } => { - *request_meta = Some(meta); - } - _ => {} - } +impl ElicitationRequest { + pub fn respond(self, result: ElicitResult) { + let _ = self.responder.send(result); + } + + pub(crate) fn new(server: String, request: ElicitRequestParams) -> (Self, oneshot::Receiver) { + let (responder, response) = oneshot::channel(); + (Self { server, request, responder }, response) } - request +} + +pub(crate) fn cancelled() -> ElicitResult { + ElicitResult::new(ElicitationAction::Cancel) } diff --git a/crates/mcp-utils/src/client/error.rs b/crates/mcp-utils/src/client/error.rs index dbfeb7329..01fb2d4e8 100644 --- a/crates/mcp-utils/src/client/error.rs +++ b/crates/mcp-utils/src/client/error.rs @@ -1,79 +1,60 @@ +use aether_auth::OAuthError; +use rmcp::service::{ClientInitializeError, ServiceError}; +use std::io; +use std::path::PathBuf; +use std::time::Duration; use thiserror::Error; #[derive(Debug, Error)] -pub enum McpError { - /// Tool not found in the registry - #[error("Tool not found: {0}")] - ToolNotFound(String), - /// The tool is model-visible and cannot be called through the deferred route. - #[error("Tool '{tool_name}' is exposed directly; call '{direct_name}' instead")] - DirectToolRequiresDirectRoute { tool_name: String, direct_name: String }, - /// Invalid tool name format (should be `server__tool`) - #[error("Invalid tool name format: {0}")] - InvalidToolNameFormat(String), - /// MCP server not found - #[error("Server not found: {0}")] - ServerNotFound(String), - /// Failed to execute tool on MCP server - #[error("Failed to execute tool {tool_name} on server {server_name}: {error}")] - ToolExecutionFailed { tool_name: String, server_name: String, error: String }, - /// Tool execution returned an error - #[error("Tool execution failed: {0}")] - ToolExecutionError(String), - /// Tool discovery failed - #[error("Tool discovery failed: {0}")] - ToolDiscoveryFailed(String), - /// Prompt not found in the registry - #[error("Prompt not found: {0}")] - PromptNotFound(String), - /// Prompt listing failed - #[error("Prompt listing failed: {0}")] - PromptListFailed(String), - /// Prompt retrieval failed - #[error("Prompt retrieval failed: {0}")] - PromptGetFailed(String), - /// A configured server conflicts with an Aether-owned instruction namespace. - #[error("MCP server name '{0}' is reserved by Aether")] - ReservedServerName(String), - /// No factory was registered for a declarative in-memory server. - #[error("In-memory MCP server '{server}' requires unregistered factory '{factory}'")] - InMemoryFactoryNotFound { server: String, factory: String }, - /// Failed to spawn a stdio server process - #[error("Failed to spawn '{command}': {reason}")] - SpawnFailed { command: String, reason: String }, - /// Server connection failed - #[error("Connection failed: {0}")] - ConnectionFailed(String), - /// Server startup failed - #[error("Server startup failed: {0}")] - ServerStartupFailed(String), - /// Transport error - #[error("Transport error: {0}")] - TransportError(String), - /// JSON serialization/deserialization error - #[error("JSON error: {0}")] - JsonError(String), - /// Generic error for other cases - #[error("{0}")] - Other(String), +pub enum ClientError { + #[error("Failed to spawn '{command}': {source}")] + Spawn { + command: String, + #[source] + source: io::Error, + }, + #[error("Failed to connect to MCP socket {}: {source}", path.display())] + SocketConnect { + path: PathBuf, + #[source] + source: io::Error, + }, + #[error("Failed to connect to MCP server '{server}': {source}")] + Connect { + server: String, + #[source] + source: Box, + }, + #[error("MCP server '{server}' requires OAuth authorization")] + AuthRequired { server: String, challenge: Option }, + #[error("MCP server '{server}' cannot be authorized interactively")] + OAuthUnavailable { server: String }, + #[error("OAuth failed for MCP server '{server}': {source}")] + Auth { + server: String, + #[source] + source: OAuthError, + }, + #[error("Authentication for MCP server '{server}' timed out after {timeout:?}")] + AuthTimedOut { server: String, timeout: Duration }, + #[error("Request to MCP server '{server}' failed: {source}")] + Request { + server: String, + #[source] + source: Box, + }, + #[error("MCP server '{server}' disconnected")] + Disconnected { server: String }, } -impl From for McpError { - fn from(error: serde_json::Error) -> Self { - McpError::JsonError(error.to_string()) +impl ClientError { + pub(crate) fn connect(server: impl Into, source: ClientInitializeError) -> Self { + Self::Connect { server: server.into(), source: Box::new(source) } } -} - -impl From for McpError { - fn from(error: std::io::Error) -> Self { - McpError::TransportError(error.to_string()) - } -} -impl From for McpError { - fn from(error: rmcp::service::ClientInitializeError) -> Self { - McpError::ConnectionFailed(error.to_string()) + pub(crate) fn request(server: impl Into, source: ServiceError) -> Self { + Self::Request { server: server.into(), source: Box::new(source) } } } -pub type Result = std::result::Result; +pub(crate) type Result = std::result::Result; diff --git a/crates/mcp-utils/src/client/handler.rs b/crates/mcp-utils/src/client/handler.rs new file mode 100644 index 000000000..ce4bd6902 --- /dev/null +++ b/crates/mcp-utils/src/client/handler.rs @@ -0,0 +1,136 @@ +use super::elicitation::{Elicitation, ElicitationRequest, cancelled}; +use super::{ClientOptions, ToolCallError}; +use rmcp::{ + ClientHandler, RoleClient, + handler::client::progress::ProgressDispatcher, + model::{ + ClientConfig, CustomNotification, ElicitRequestParams, ElicitResult, ElicitationAction, ErrorData, + InputRequest, InputRequests, InputResponses, ProgressNotificationParam, RequestMetaObject, + }, + service::{MaybeSendFuture, NotificationContext, RequestContext}, +}; +use std::future::{Future, ready}; +use tokio::sync::{mpsc, watch}; + +pub(crate) struct Handler { + pub(crate) progress: ProgressDispatcher, + info: ClientConfig, + server: String, + tx: Option>, + tool_list_changed: watch::Sender<()>, +} + +impl Handler { + pub(crate) fn new(server: String, options: &ClientOptions) -> Self { + Self { + info: options.client_config(), + server, + progress: ProgressDispatcher::new(), + tx: options.elicitations().cloned(), + tool_list_changed: watch::Sender::new(()), + } + } + + pub(crate) fn tool_list_changes(&self) -> watch::Receiver<()> { + self.tool_list_changed.subscribe() + } + + pub(crate) fn server(&self) -> &str { + &self.server + } + + pub(crate) async fn elicit_inputs(&self, requests: InputRequests) -> Result<(InputResponses, bool), ToolCallError> { + let mut responses = InputResponses::new(); + let mut cancelled = false; + for (key, request) in requests { + let InputRequest::Elicitation(elicitation_request) = request else { + return Err(ToolCallError::UnsupportedInput); + }; + let extension_meta = elicitation_request.extensions.get::().cloned(); + let result = self.elicit(elicitation_request.params, extension_meta).await; + cancelled |= result.action == ElicitationAction::Cancel; + let response = serde_json::to_value(&result).expect("ElicitResult serializes to JSON"); + responses.insert(key, response); + } + Ok((responses, cancelled)) + } + + async fn elicit(&self, request: ElicitRequestParams, meta: Option) -> ElicitResult { + let Some(events) = &self.tx else { + return cancelled(); + }; + let (elicitation, response) = ElicitationRequest::new(self.server.clone(), with_meta(request, meta)); + if events.send(Elicitation::Request(Box::new(elicitation))).await.is_err() { + return cancelled(); + } + response.await.unwrap_or_else(|_| cancelled()) + } +} + +impl ClientHandler for Handler { + fn get_info(&self) -> ClientConfig { + self.info.clone() + } + + async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext) -> () { + self.progress.handle_notification(params).await; + } + + async fn create_elicitation( + &self, + request: ElicitRequestParams, + context: RequestContext, + ) -> Result { + let meta = (!context.meta.is_empty()).then_some(context.meta); + Ok(self.elicit(request, meta).await) + } + + async fn on_custom_notification( + &self, + notification: CustomNotification, + _context: NotificationContext, + ) { + if notification.method != "notifications/elicitation/complete" { + return; + } + let params: Option = + notification.params.and_then(|params| serde_json::from_value(params).ok()); + + let Some(params) = params else { + tracing::warn!("Ignoring malformed MCP elicitation completion notification"); + return; + }; + + if let Some(events) = &self.tx { + let completion = Elicitation::Complete { server: self.server.clone(), id: params.elicitation_id }; + let _ = events.send(completion).await; + } + } + + fn on_tool_list_changed( + &self, + _context: NotificationContext, + ) -> impl Future + MaybeSendFuture + '_ { + self.tool_list_changed.send_replace(()); + ready(()) + } +} + +#[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct ElicitationCompleteParams { + elicitation_id: String, +} + +fn with_meta(mut request: ElicitRequestParams, meta: Option) -> ElicitRequestParams { + if let Some(meta) = meta { + match &mut request { + ElicitRequestParams::FormElicitationParams { meta: request_meta, .. } + | ElicitRequestParams::UrlElicitationParams { meta: request_meta, .. } => { + *request_meta = Some(meta); + } + _ => {} + } + } + request +} diff --git a/crates/mcp-utils/src/client/manager.rs b/crates/mcp-utils/src/client/manager.rs deleted file mode 100644 index 5ad502eff..000000000 --- a/crates/mcp-utils/src/client/manager.rs +++ /dev/null @@ -1,1086 +0,0 @@ -use llm::ToolDefinition; - -use super::{ - McpError, McpSnapshot, Result, - config::{McpHttpConfig, ToolExposure}, - connection::{ - ConnectConfig, McpConnectAttempt, McpConnectOutcome, McpServerConnection, Tool, authenticate_http, - connect_server, - }, - mcp_client::client_capabilities, - naming::{create_namespaced_tool_name, split_on_server_name}, - tool_catalog::{ServerCatalogEntry, ToolCatalog}, - tool_filter::ToolFilter, -}; -use aether_auth::{OAuthCredentialStorage, OAuthHandler}; -use futures::future::join_all; -use rmcp::{ - Peer, RoleClient, RoleServer, - model::{ClientCapabilities, ClientConfig, ElicitRequestParams, ElicitResult, Implementation, Tool as RmcpTool}, - service::DynService, -}; -use std::collections::{BTreeMap, HashMap}; -use std::future::Future; -use std::num::NonZeroU16; -use std::path::PathBuf; -use std::sync::{Arc, atomic::AtomicU64}; -use tokio::sync::{mpsc, oneshot, watch}; -use tokio::task::JoinHandle; - -use utils::mcp_status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry}; - -pub type OAuthHandlerFactory = Arc Result> + Send + Sync>; - -pub struct ToolListChangedRequest { - server: String, - generation: u64, - peer: Peer, -} - -pub struct ToolListRefresh { - server: String, - generation: u64, - result: Result>, -} - -impl ToolListChangedRequest { - pub(crate) fn new(server: String, generation: u64, peer: Peer) -> Self { - Self { server, generation, peer } - } - - pub async fn refresh(self) -> ToolListRefresh { - let result = self - .peer - .list_all_tools() - .await - .map_err(|error| McpError::ToolDiscoveryFailed(format!("Failed to refresh tools: {error}"))); - ToolListRefresh { server: self.server, generation: self.generation, result } - } -} - -pub struct RuntimeMcpServer { - pub name: String, - pub transport: RuntimeMcpTransport, - pub tool_exposure: ToolExposure, -} - -pub enum RuntimeMcpTransport { - Stdio { command: String, args: Vec, env: HashMap }, - Http(McpHttpConfig), - InMemory { server: Box> }, -} - -impl RuntimeMcpServer { - pub fn new(name: impl Into, transport: RuntimeMcpTransport, tool_exposure: ToolExposure) -> Self { - Self { name: name.into(), transport, tool_exposure } - } - - pub fn with_exposure(mut self, exposure: ToolExposure) -> Self { - self.tool_exposure = exposure; - self - } -} - -/// Context passed to an `OAuthHandlerFactory` so the constructed handler can -/// dispatch user-facing prompts back to the host through the MCP event channel. -#[derive(Clone)] -pub struct OAuthHandlerContext { - pub server_name: String, - pub callback_port: Option, - pub tx: mpsc::Sender, -} - -#[derive(Debug)] -pub struct ElicitationRequest { - pub server_name: String, - pub request: ElicitRequestParams, - pub response_sender: oneshot::Sender, -} - -/// Events emitted by MCP clients that require attention from the host -/// (e.g. the relay or TUI). Flows through a single channel from `McpManager` -/// to the consumer. -#[derive(Debug)] -pub enum McpClientEvent { - Elicitation(Box), - ElicitationComplete { server_name: String, elicitation_id: String }, - ServerStatusesChanged(Vec), - AuthenticationFailed { server: String, error: String }, - ConnectionReady(McpConnectionDetails), -} - -pub type McpConnectionDetails = Arc; - -/// Manages connections to multiple MCP servers and their tools -pub struct McpManager { - servers: HashMap, - catalog: ToolCatalog, - tool_filter: ToolFilter, - client_info: ClientConfig, - event_sender: mpsc::Sender, - root_dir: PathBuf, - oauth_handler_factory: Option, - oauth_credential_store: Option>, - snapshot_sender: Option>>, - tool_refresh_sender: mpsc::Sender, - tool_refresh_receiver: Option>, - next_connection_generation: Arc, - progressive_discovery_instructions: Option, -} - -impl McpManager { - pub fn new(event_sender: mpsc::Sender, oauth_handler_factory: Option) -> Self { - let (tool_refresh_sender, tool_refresh_receiver) = mpsc::channel(32); - Self { - servers: HashMap::new(), - catalog: ToolCatalog::new(), - tool_filter: ToolFilter::default(), - client_info: ClientConfig::new(client_capabilities(), Implementation::new("aether", "0.1.0")), - event_sender, - root_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), - oauth_handler_factory, - oauth_credential_store: None, - snapshot_sender: None, - tool_refresh_sender, - tool_refresh_receiver: Some(tool_refresh_receiver), - next_connection_generation: Arc::new(AtomicU64::new(1)), - progressive_discovery_instructions: None, - } - } - - pub fn take_tool_refresh_receiver(&mut self) -> mpsc::Receiver { - self.tool_refresh_receiver.take().expect("tool refresh receiver can only be taken once") - } - - pub fn with_client_capabilities(mut self, capabilities: ClientCapabilities) -> Self { - self.client_info.capabilities = capabilities; - self - } - - pub fn with_progressive_discovery_instructions(mut self, instructions: impl Into) -> Self { - self.progressive_discovery_instructions = Some(instructions.into()); - self - } - - pub fn with_snapshot_sender(mut self, sender: watch::Sender>) -> Self { - self.snapshot_sender = Some(sender); - self.publish_snapshot(); - self - } - - pub fn with_oauth_credential_store(mut self, store: Arc) -> Self { - self.oauth_credential_store = Some(store); - self - } - - pub fn with_root_dir(mut self, root_dir: impl Into) -> Self { - self.root_dir = root_dir.into(); - self - } - - pub fn with_tool_filter(mut self, filter: ToolFilter) -> Self { - self.tool_filter = filter; - self - } - - pub fn catalog(&self) -> &ToolCatalog { - &self.catalog - } - - pub async fn register_pending(&mut self, servers: Vec) -> Result> { - for server in &servers { - self.register_record(&server.name, ServerState::Connecting, None, server.tool_exposure.clone()); - } - - self.publish_snapshot(); - self.emit_server_statuses_changed().await; - Ok(servers) - } - - pub fn connect_pending_task( - &self, - server: RuntimeMcpServer, - ) -> impl Future + Send + 'static { - let ctx = self.connect_config(); - async move { connect_server(server, &ctx).await } - } - - pub async fn add_mcps(&mut self, servers: Vec) -> Result<()> { - let pending = self.register_pending(servers).await?; - let ctx = self.connect_config(); - let attempts = join_all(pending.into_iter().map(|server| connect_server(server, &ctx))).await; - for attempt in attempts { - self.apply_connection_attempt(attempt).await; - } - Ok(()) - } - - pub fn tool_definitions(&self) -> Vec { - self.catalog.tools().model_visible.into_iter().map(|tool| tool.definition().clone()).collect() - } - - pub fn server_instructions(&self) -> BTreeMap { - self.catalog.model_instructions() - } - - pub fn server_statuses(&self) -> Vec { - self.catalog.server_statuses() - } - - pub async fn authenticate_server_task( - &mut self, - name: &str, - ) -> Result + Send + 'static> { - let record = self - .servers - .get(name) - .ok_or_else(|| McpError::ConnectionFailed(format!("server '{name}' is not OAuth-authenticatable")))?; - if !record.can_authenticate() { - return Err(McpError::ConnectionFailed(format!("server '{name}' is not OAuth-authenticatable"))); - } - if self.oauth_handler_factory.is_none() { - return Err(McpError::ConnectionFailed(format!("No OAuth handler factory available for '{name}'"))); - } - - let name = name.to_string(); - let config = record.reauth_config.clone().expect("checked above"); - let challenge = record.oauth_challenge.clone(); - let ctx = self.connect_config(); - - self.set_state(&name, ServerState::Authenticating); - self.emit_server_statuses_changed().await; - - Ok(async move { authenticate_http(name, config, challenge, ctx).await }) - } - - pub async fn apply_tool_list_refresh(&mut self, refresh: ToolListRefresh) { - let ToolListRefresh { server, generation, result } = refresh; - let Some(record) = self.servers.get(&server) else { - return; - }; - let Some(connection) = record.connection() else { - return; - }; - if connection.generation() != generation { - tracing::debug!(server = %server, generation, "Ignoring stale MCP tool refresh"); - return; - } - let tools = match result { - Ok(tools) => tools, - Err(error) => { - tracing::warn!(server = %server, %error, "Failed to refresh MCP tools; retaining previous catalog"); - return; - } - }; - if let Err(error) = self.replace_catalog_tools(&server, &tools) { - tracing::warn!(server = %server, %error, "Failed to apply refreshed MCP tools; retaining previous catalog"); - return; - } - self.emit_server_statuses_changed().await; - } - - pub async fn apply_connection_attempt(&mut self, attempt: McpConnectAttempt) { - let McpConnectAttempt { name, outcome } = attempt; - match outcome { - McpConnectOutcome::Connected { conn, reauth_config } => { - match self.register_connection(&name, conn, reauth_config).await { - Ok(()) => { - self.emit_server_statuses_changed().await; - } - Err(error) => self.apply_authentication_failure(name, error.to_string()).await, - } - } - McpConnectOutcome::NeedsOAuth { config, challenge, error } => { - tracing::warn!("Server '{name}' needs OAuth: {error}"); - if let Some(record) = self.servers.get_mut(&name) { - record.reauth_config = Some(config); - record.oauth_challenge = challenge; - } - self.set_state(&name, ServerState::NeedsOAuth); - self.emit_server_statuses_changed().await; - } - McpConnectOutcome::Failed { error } => { - self.apply_authentication_failure(name, error.to_string()).await; - } - } - } - - /// List all prompts from all connected MCP servers with namespacing - pub async fn list_prompts(&self) -> Result> { - let futures: Vec<_> = self - .servers - .iter() - .filter_map(|(server_name, record)| { - let conn = record.connection()?; - conn.client.peer_info()?.capabilities.prompts.as_ref()?; - let server_name = server_name.clone(); - let client = conn.client.clone(); - Some(async move { - let prompts = client.list_all_prompts().await.map_err(|e| { - McpError::PromptListFailed(format!("Failed to list prompts for {server_name}: {e}")) - })?; - - let namespaced_prompts: Vec = prompts - .into_iter() - .map(|prompt| { - let namespaced_name = create_namespaced_tool_name(&server_name, &prompt.name); - rmcp::model::Prompt::new(namespaced_name, prompt.description, prompt.arguments) - }) - .collect(); - - Ok::<_, McpError>(namespaced_prompts) - }) - }) - .collect(); - - let results = join_all(futures).await; - let mut all_prompts = Vec::new(); - for result in results { - all_prompts.extend(result?); - } - - Ok(all_prompts) - } - - /// Get a specific prompt by namespaced name - pub async fn get_prompt( - &self, - namespaced_prompt_name: &str, - arguments: Option>, - ) -> Result { - let (server_name, prompt_name) = split_on_server_name(namespaced_prompt_name) - .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_prompt_name.to_string()))?; - - let server_conn = - self.connection_for(server_name).ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?; - - let mut request = rmcp::model::GetPromptRequestParams::new(prompt_name); - if let Some(args) = arguments { - request = request.with_arguments(args); - } - - server_conn.client.get_prompt(request).await.map_err(|e| { - McpError::PromptGetFailed(format!("Failed to get prompt '{prompt_name}' from {server_name}: {e}")) - }) - } - - /// Shutdown all servers and wait for their tasks to complete - pub async fn shutdown(&mut self) { - let servers: Vec<(String, ServerRecord)> = self.servers.drain().collect(); - self.catalog = ToolCatalog::new(); - self.publish_snapshot(); - - for (server_name, record) in servers { - if let Some(conn) = record.into_connection() - && let Some(handle) = conn.server_task - { - drop(conn.client); - await_server_shutdown(&server_name, handle).await; - } - } - } - - /// Shutdown a specific server by name - pub async fn shutdown_server(&mut self, server_name: &str) -> Result<()> { - self.catalog.remove_server(server_name); - self.refresh_progressive_discovery_instructions(); - let record = self.servers.remove(server_name); - self.publish_snapshot(); - if let Some(record) = record - && let Some(conn) = record.into_connection() - && let Some(handle) = conn.server_task - { - drop(conn.client); - await_server_shutdown(server_name, handle).await; - } - - Ok(()) - } - - async fn emit_server_statuses_changed(&self) { - self.emit_event(McpClientEvent::ServerStatusesChanged(self.server_statuses())).await; - } - - fn refresh_progressive_discovery_instructions(&mut self) { - let instructions = (!self.catalog.discoverable_deferred_servers().is_empty()) - .then(|| self.progressive_discovery_instructions.clone()) - .flatten(); - self.catalog.set_progressive_discovery_instructions(instructions); - } - - pub async fn emit_connection_ready(&self) { - self.emit_event(McpClientEvent::ConnectionReady(self.snapshot())).await; - } - - async fn emit_authentication_failed(&self, server: String, error: String) { - self.emit_event(McpClientEvent::AuthenticationFailed { server, error }).await; - } - - async fn emit_event(&self, event: McpClientEvent) { - if let Err(e) = self.event_sender.send(event).await { - tracing::warn!("Failed to emit MCP client event: {e}"); - } - } - - fn connect_config(&self) -> Arc { - Arc::new(ConnectConfig { - client_info: self.client_info.clone(), - event_sender: self.event_sender.clone(), - tool_refresh_sender: self.tool_refresh_sender.clone(), - next_connection_generation: Arc::clone(&self.next_connection_generation), - root_dir: self.root_dir.clone(), - oauth_handler_factory: self.oauth_handler_factory.clone(), - oauth_credential_store: self.oauth_credential_store.clone(), - }) - } - - async fn register_connection( - &mut self, - name: &str, - conn: McpServerConnection, - reauth_config: Option, - ) -> Result<()> { - let tools = conn - .list_tools() - .await - .map_err(|e| McpError::ToolDiscoveryFailed(format!("Failed to list tools for {name}: {e}")))?; - let exposure = - self.catalog.server(name).ok_or_else(|| McpError::ServerNotFound(name.to_string()))?.exposure().clone(); - - let auth_capability = { - let record = self.servers.get_mut(name).expect("record checked above"); - record.reauth_config = reauth_config.or_else(|| record.reauth_config.clone()); - record.oauth_challenge = None; - record.auth_capability() - }; - let description = conn - .client - .peer_info() - .and_then(|info| info.server_info.as_ref()?.description.clone()) - .filter(|description| !description.is_empty()) - .unwrap_or_else(|| name.to_string()); - let instructions = conn.instructions.clone(); - let catalog_tools = tools.iter().map(Tool::from).collect::>(); - let entry = ServerCatalogEntry::from_tools( - name.to_string(), - description, - instructions, - McpServerStatus::Connected { tool_count: catalog_tools.len() }, - auth_capability, - exposure.clone(), - &catalog_tools, - &self.tool_filter, - ); - - self.servers.get_mut(name).expect("record checked above").state = ServerState::Connected { connection: conn }; - self.catalog.upsert_server(entry); - self.refresh_progressive_discovery_instructions(); - self.publish_snapshot(); - Ok(()) - } - - fn replace_catalog_tools(&mut self, name: &str, tools: &[RmcpTool]) -> Result<()> { - let existing = self.catalog.server(name).cloned().ok_or_else(|| McpError::ServerNotFound(name.to_string()))?; - let catalog_tools = tools.iter().map(Tool::from).collect::>(); - let entry = ServerCatalogEntry::from_tools( - name.to_string(), - existing.description().to_string(), - existing.instructions().map(str::to_string), - McpServerStatus::Connected { tool_count: catalog_tools.len() }, - existing.auth_capability(), - existing.exposure().clone(), - &catalog_tools, - &self.tool_filter, - ); - self.catalog.upsert_server(entry); - self.refresh_progressive_discovery_instructions(); - self.publish_snapshot(); - Ok(()) - } - - async fn apply_authentication_failure(&mut self, name: String, error: String) { - self.set_state(&name, ServerState::Failed { error: error.clone() }); - self.emit_server_statuses_changed().await; - self.emit_authentication_failed(name, error).await; - } - - fn set_state(&mut self, name: &str, state: ServerState) { - let status = McpServerStatus::from(&state); - match self.servers.get_mut(name) { - Some(record) => record.state = state, - None => { - self.servers.insert(name.to_string(), ServerRecord::new(state, None)); - } - } - let auth = self.servers.get(name).map_or(McpServerAuthCapability::Unavailable, ServerRecord::auth_capability); - let entry = self - .catalog - .server(name) - .cloned() - .unwrap_or_else(|| ServerCatalogEntry::pending(name, ToolExposure::ModelVisible)) - .with_status(status, auth); - self.catalog.upsert_server(entry); - self.refresh_progressive_discovery_instructions(); - self.publish_snapshot(); - } - - fn register_record( - &mut self, - name: &str, - state: ServerState, - reauth_config: Option, - exposure: ToolExposure, - ) { - let status = McpServerStatus::from(&state); - let auth_capability = - if reauth_config.is_some() { McpServerAuthCapability::OAuth } else { McpServerAuthCapability::Unavailable }; - self.servers.insert(name.to_string(), ServerRecord::new(state, reauth_config)); - self.catalog.upsert_server(ServerCatalogEntry::pending(name, exposure).with_status(status, auth_capability)); - } - - pub fn snapshot(&self) -> Arc { - let clients = self - .servers - .iter() - .filter_map(|(name, record)| record.connection().map(|conn| (name.clone(), conn.client.clone()))) - .collect(); - Arc::new(McpSnapshot::new(Arc::new(self.catalog.clone()), Arc::new(clients))) - } - - fn publish_snapshot(&self) { - if let Some(sender) = &self.snapshot_sender { - sender.send_replace(self.snapshot()); - } - } - - fn connection_for(&self, server_name: &str) -> Option<&McpServerConnection> { - self.servers.get(server_name).and_then(ServerRecord::connection) - } -} - -impl Drop for McpManager { - fn drop(&mut self) { - let servers: Vec<(String, ServerRecord)> = self.servers.drain().collect(); - for (server_name, record) in servers { - if let Some(conn) = record.into_connection() - && let Some(handle) = conn.server_task - { - handle.abort(); - tracing::warn!("Server '{server_name}' task aborted during cleanup"); - } - } - } -} - -/// Internal record holding all mutable state for a single MCP server. -struct ServerRecord { - state: ServerState, - reauth_config: Option, - oauth_challenge: Option, -} - -enum ServerState { - Connecting, - Connected { connection: McpServerConnection }, - Authenticating, - Failed { error: String }, - NeedsOAuth, -} - -impl From<&ServerState> for McpServerStatus { - fn from(state: &ServerState) -> Self { - match state { - ServerState::Connecting => Self::Connecting, - ServerState::Connected { .. } => Self::Connected { tool_count: 0 }, - ServerState::Authenticating => Self::Authenticating, - ServerState::Failed { error } => Self::Failed { error: error.clone() }, - ServerState::NeedsOAuth => Self::NeedsOAuth, - } - } -} - -impl ServerRecord { - fn new(state: ServerState, reauth_config: Option) -> Self { - Self { state, reauth_config, oauth_challenge: None } - } - - fn connection(&self) -> Option<&McpServerConnection> { - match &self.state { - ServerState::Connected { connection, .. } => Some(connection), - ServerState::Connecting - | ServerState::Authenticating - | ServerState::Failed { .. } - | ServerState::NeedsOAuth => None, - } - } - - fn into_connection(self) -> Option { - match self.state { - ServerState::Connected { connection, .. } => Some(connection), - ServerState::Connecting - | ServerState::Authenticating - | ServerState::Failed { .. } - | ServerState::NeedsOAuth => None, - } - } - - fn auth_capability(&self) -> McpServerAuthCapability { - if self.reauth_config.is_some() { McpServerAuthCapability::OAuth } else { McpServerAuthCapability::Unavailable } - } - - fn can_authenticate(&self) -> bool { - self.reauth_config.is_some() - } -} - -/// Awaits `handle` for up to 5 seconds, logging whether the server shut down -/// gracefully, panicked, or timed out. Used during manager teardown. -async fn await_server_shutdown(server_name: &str, handle: JoinHandle<()>) { - let Ok(task_result) = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await else { - tracing::warn!("Server '{server_name}' shutdown timed out"); - return; - }; - match task_result { - Ok(()) => tracing::info!("Server '{server_name}' shut down gracefully"), - Err(e) => tracing::warn!("Server '{server_name}' task panicked: {e:?}"), - } -} - -#[cfg(test)] -mod tests { - use super::{ - McpClientEvent, McpManager, McpServerStatus, RuntimeMcpServer as McpServer, - RuntimeMcpTransport as McpTransport, ServerState, ToolListRefresh, - }; - use crate::client::config::{McpHttpConfig, ToolExposure}; - use crate::client::connection::{McpConnectAttempt, McpConnectOutcome}; - use crate::client::{McpSnapshot, OAuthHandlerFactory, ToolRoute}; - use aether_auth::{OAuthError, OAuthHandler}; - use futures::future::BoxFuture; - use rmcp::{ - Json, RoleServer, ServerHandler, - handler::server::{router::tool::ToolRouter, wrapper::Parameters}, - model::{Implementation, ServerCapabilities, ServerConfig, Tool as RmcpTool}, - service::DynService, - tool, tool_handler, tool_router, - transport::streamable_http_client::StreamableHttpClientTransportConfig, - }; - use schemars::JsonSchema; - use serde::{Deserialize, Serialize}; - use std::{ - io, - sync::{Arc, Mutex}, - }; - use tokio::sync::{mpsc, watch}; - use utils::mcp_status::McpServerAuthCapability; - - #[derive(Clone)] - struct TestServer { - tool_router: ToolRouter, - } - - #[allow(clippy::unused_async_trait_impl)] - #[tool_handler(router = self.tool_router)] - impl ServerHandler for TestServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("test-server", "0.1.0").with_description("Test MCP server")) - .with_instructions("Test server instructions") - } - } - - impl Default for TestServer { - fn default() -> Self { - Self { tool_router: Self::tool_router() } - } - } - - #[derive(Debug, Deserialize, Serialize, JsonSchema)] - struct EchoRequest { - value: String, - } - - #[derive(Debug, Deserialize, Serialize, JsonSchema)] - struct EchoResult { - value: String, - } - - #[tool_router] - impl TestServer { - fn into_dyn(self) -> Box> { - Box::new(self) - } - - #[tool(description = "Returns the provided value", annotations(read_only_hint = true, open_world_hint = false))] - async fn echo(&self, request: Parameters) -> Json { - let Parameters(EchoRequest { value }) = request; - Json(EchoResult { value }) - } - } - - #[derive(Clone)] - struct SharedWriter(Arc>>); - - impl io::Write for SharedWriter { - fn write(&mut self, buf: &[u8]) -> io::Result { - self.0.lock().unwrap().extend_from_slice(buf); - Ok(buf.len()) - } - - fn flush(&mut self) -> io::Result<()> { - Ok(()) - } - } - - struct TestOAuthHandler; - - impl OAuthHandler for TestOAuthHandler { - fn redirect_uri(&self) -> &'static str { - "http://127.0.0.1:0/oauth2callback" - } - - fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { - Box::pin(async { Err(OAuthError::UserCancelled) }) - } - } - - fn test_oauth_handler_factory() -> OAuthHandlerFactory { - Arc::new(|_ctx| Ok(Arc::new(TestOAuthHandler))) - } - - fn http_config(uri: &str) -> McpHttpConfig { - StreamableHttpClientTransportConfig::with_uri(uri).into() - } - - #[tokio::test] - async fn authenticate_server_task_rejects_record_without_reauth_config() { - let (event_sender, _event_receiver) = mpsc::channel(1); - let mut manager = McpManager::new(event_sender, Some(test_oauth_handler_factory())); - manager.register_record("public", ServerState::Connecting, None, ToolExposure::ModelVisible); - - let error = match manager.authenticate_server_task("public").await { - Ok(_) => panic!("non-OAuth server should be rejected"), - Err(error) => error.to_string(), - }; - assert!(error.contains("not OAuth-authenticatable")); - } - - #[tokio::test] - async fn authenticate_server_task_marks_server_authenticating_and_emits_status() { - let (event_sender, mut event_receiver) = mpsc::channel(2); - let mut manager = McpManager::new(event_sender, Some(test_oauth_handler_factory())); - manager.register_record( - "remote", - ServerState::NeedsOAuth, - Some(http_config("http://localhost:19999/mcp")), - ToolExposure::ModelVisible, - ); - - let _task = manager.authenticate_server_task("remote").await.expect("auth should start"); - - assert!(matches!(manager.server_statuses()[0].status, McpServerStatus::Authenticating)); - let event = event_receiver.recv().await.expect("status change event"); - let McpClientEvent::ServerStatusesChanged(servers) = event else { - panic!("expected ServerStatusesChanged"); - }; - let status = servers.iter().find(|entry| entry.name == "remote").expect("remote status"); - assert!(matches!(status.status, McpServerStatus::Authenticating)); - assert_eq!(status.auth_capability, McpServerAuthCapability::OAuth); - } - - #[tokio::test] - async fn apply_connection_attempt_failure_allows_retry() { - let (event_sender, mut event_receiver) = mpsc::channel(2); - let mut manager = McpManager::new(event_sender, Some(test_oauth_handler_factory())); - manager.register_record( - "remote", - ServerState::NeedsOAuth, - Some(http_config("http://localhost:19999/mcp")), - ToolExposure::ModelVisible, - ); - let _task = manager.authenticate_server_task("remote").await.expect("auth should start"); - let _authenticating_event = event_receiver.recv().await.expect("authenticating status change event"); - - manager - .apply_connection_attempt(McpConnectAttempt { - name: "remote".to_string(), - outcome: McpConnectOutcome::Failed { - error: crate::client::McpError::ConnectionFailed("boom".to_string()), - }, - }) - .await; - - let event = event_receiver.recv().await.expect("status change event"); - let McpClientEvent::ServerStatusesChanged(servers) = event else { - panic!("expected ServerStatusesChanged"); - }; - let auth_event = event_receiver.recv().await.expect("authentication failure event"); - let McpClientEvent::AuthenticationFailed { server, error } = auth_event else { - panic!("expected AuthenticationFailed"); - }; - assert_eq!(server, "remote"); - assert!(error.contains("boom")); - - let status = servers.iter().find(|entry| entry.name == "remote").expect("remote status"); - assert_eq!(status.auth_capability, McpServerAuthCapability::OAuth); - assert!(matches!(status.status, McpServerStatus::Failed { ref error } if error.contains("boom"))); - assert!(manager.authenticate_server_task("remote").await.is_ok()); - } - - #[test] - fn status_entries_are_derived_from_reauth_config() { - let (event_sender, _event_receiver) = mpsc::channel(1); - let mut manager = McpManager::new(event_sender, Some(test_oauth_handler_factory())); - - manager.register_record( - "with-oauth", - ServerState::Connecting, - Some(http_config("http://localhost/mcp")), - ToolExposure::ModelVisible, - ); - manager.register_record("without-oauth", ServerState::Connecting, None, ToolExposure::ModelVisible); - manager.register_record( - "needs-oauth", - ServerState::NeedsOAuth, - Some(http_config("http://localhost/mcp2")), - ToolExposure::ModelVisible, - ); - - let statuses = manager.server_statuses(); - let with_oauth = statuses.iter().find(|s| s.name == "with-oauth").unwrap(); - let without_oauth = statuses.iter().find(|s| s.name == "without-oauth").unwrap(); - let needs_oauth = statuses.iter().find(|s| s.name == "needs-oauth").unwrap(); - - assert_eq!(with_oauth.auth_capability, McpServerAuthCapability::OAuth); - assert_eq!(without_oauth.auth_capability, McpServerAuthCapability::Unavailable); - assert_eq!(needs_oauth.auth_capability, McpServerAuthCapability::OAuth); - } - - #[tokio::test] - async fn register_pending_marks_every_server_connecting_and_emits_status() { - let (event_sender, mut event_receiver) = mpsc::channel(32); - let mut manager = McpManager::new(event_sender, None); - - let servers = vec![ - McpServer::new( - "alpha", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - ), - McpServer::new( - "beta", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::deferred_all(), - ), - ]; - - let returned = manager.register_pending(servers).await.unwrap(); - assert_eq!(returned.iter().map(|s| s.name.as_str()).collect::>(), vec!["alpha", "beta"]); - - let statuses = manager.server_statuses(); - assert_eq!(statuses.len(), 2); - assert!(matches!(statuses.iter().find(|s| s.name == "alpha").unwrap().status, McpServerStatus::Connecting)); - assert!(matches!(statuses.iter().find(|s| s.name == "beta").unwrap().status, McpServerStatus::Connecting)); - assert!(statuses.iter().find(|s| s.name == "beta").unwrap().deferred_tools); - - let event = event_receiver.try_recv().expect("expected initial ServerStatusesChanged emission"); - let McpClientEvent::ServerStatusesChanged(emitted) = event else { - panic!("expected ServerStatusesChanged, got {event:?}"); - }; - assert_eq!(emitted.iter().map(|s| s.name.as_str()).collect::>(), vec!["alpha", "beta"]); - } - - #[tokio::test] - async fn server_statuses_mark_model_visible_and_deferred_servers() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let mut manager = McpManager::new(event_sender, None); - manager - .add_mcps(vec![ - McpServer::new( - "direct", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - ), - McpServer::new( - "math", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::deferred_all(), - ), - ]) - .await - .unwrap(); - - let statuses = manager.server_statuses(); - assert_eq!(statuses.iter().map(|status| status.name.as_str()).collect::>(), vec!["direct", "math"]); - assert!(!statuses.iter().find(|status| status.name == "direct").unwrap().deferred_tools); - assert!(statuses.iter().find(|status| status.name == "math").unwrap().deferred_tools); - } - - #[tokio::test] - async fn tool_definitions_drop_when_a_server_shuts_down() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let mut manager = McpManager::new(event_sender, None); - manager - .add_mcps(vec![ - McpServer::new( - "git", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - ), - McpServer::new( - "github", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - ), - ]) - .await - .unwrap(); - - let names = - |manager: &McpManager| manager.tool_definitions().into_iter().map(|tool| tool.name).collect::>(); - assert!(names(&manager).contains(&"git__echo".to_string())); - assert!(names(&manager).contains(&"github__echo".to_string())); - - manager.shutdown_server("git").await.unwrap(); - - assert!(!names(&manager).iter().any(|name| name.starts_with("git__"))); - assert!(names(&manager).contains(&"github__echo".to_string())); - } - - #[tokio::test] - async fn server_removal_publishes_before_stale_connection_cleanup() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let (snapshot_sender, mut snapshots) = watch::channel(Arc::new(McpSnapshot::default())); - let mut manager = McpManager::new(event_sender, None).with_snapshot_sender(snapshot_sender); - manager - .add_mcps(vec![McpServer::new( - "test", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - )]) - .await - .unwrap(); - let connected = snapshots.borrow().clone(); - assert_eq!(connected.tool_definitions()[0].name, "test__echo"); - - manager.shutdown_server("test").await.unwrap(); - snapshots.changed().await.unwrap(); - let removed = snapshots.borrow().clone(); - - assert_eq!(removed.tool_definitions().len(), 0); - assert!( - removed - .resolve(ToolRoute::ModelVisible { namespaced_name: "test__echo".to_string() }, serde_json::Map::new(),) - .is_err() - ); - assert_eq!(connected.tool_definitions()[0].name, "test__echo"); - } - - #[tokio::test] - async fn failed_tool_refresh_preserves_last_healthy_snapshot() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let (snapshot_sender, snapshots) = watch::channel(Arc::new(McpSnapshot::default())); - let mut manager = McpManager::new(event_sender, None).with_snapshot_sender(snapshot_sender); - manager - .add_mcps(vec![McpServer::new( - "test", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - )]) - .await - .unwrap(); - let healthy = snapshots.borrow().clone(); - let generation = manager.connection_for("test").unwrap().generation(); - - manager - .apply_tool_list_refresh(ToolListRefresh { - server: "test".to_string(), - generation, - result: Err(crate::client::McpError::ToolDiscoveryFailed("boom".to_string())), - }) - .await; - - assert!(Arc::ptr_eq(&healthy, &snapshots.borrow())); - assert_eq!(snapshots.borrow().tool_definitions()[0].name, "test__echo"); - } - - #[tokio::test] - async fn stale_tool_refresh_is_ignored_after_connection_generation_changes() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let (snapshot_sender, snapshots) = watch::channel(Arc::new(McpSnapshot::default())); - let mut manager = McpManager::new(event_sender, None).with_snapshot_sender(snapshot_sender); - manager - .add_mcps(vec![McpServer::new( - "test", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - )]) - .await - .unwrap(); - let healthy = snapshots.borrow().clone(); - let stale_generation = manager.connection_for("test").unwrap().generation() + 1; - let added = RmcpTool::new("stale", "stale", Arc::new(serde_json::Map::new())); - - manager - .apply_tool_list_refresh(ToolListRefresh { - server: "test".to_string(), - generation: stale_generation, - result: Ok(vec![added]), - }) - .await; - - assert!(Arc::ptr_eq(&healthy, &snapshots.borrow())); - assert!(!snapshots.borrow().tool_definitions().iter().any(|tool| tool.name == "test__stale")); - } - - #[tokio::test] - async fn tool_definitions_preserve_annotations() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let mut manager = McpManager::new(event_sender, None); - manager - .add_mcps(vec![McpServer::new( - "test", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - )]) - .await - .unwrap(); - - let tools = manager.tool_definitions(); - let echo = tools.iter().find(|tool| tool.name == "test__echo").expect("echo tool"); - let annotations = echo.annotations.as_ref().expect("annotations should be preserved"); - assert_eq!(annotations.read_only_hint, Some(true)); - assert_eq!(annotations.open_world_hint, Some(false)); - } - - #[tokio::test] - async fn drop_logs_cleanup_abort_with_tracing() { - let (event_sender, _event_receiver) = mpsc::channel(32); - let mut manager = McpManager::new(event_sender, None); - manager - .add_mcps(vec![McpServer::new( - "test", - McpTransport::InMemory { server: TestServer::default().into_dyn() }, - ToolExposure::ModelVisible, - )]) - .await - .unwrap(); - - let output = Arc::new(Mutex::new(Vec::new())); - let subscriber = tracing_subscriber::fmt() - .with_ansi(false) - .without_time() - .with_writer({ - let output = Arc::clone(&output); - move || SharedWriter(Arc::clone(&output)) - }) - .finish(); - - tracing::subscriber::with_default(subscriber, || { - drop(manager); - }); - - let logs = String::from_utf8(output.lock().unwrap().clone()).unwrap(); - assert!(logs.contains("Server 'test' task aborted during cleanup")); - } -} diff --git a/crates/mcp-utils/src/client/mcp_client.rs b/crates/mcp-utils/src/client/mcp_client.rs index 03e3ddc6e..b91188db0 100644 --- a/crates/mcp-utils/src/client/mcp_client.rs +++ b/crates/mcp-utils/src/client/mcp_client.rs @@ -1,241 +1,165 @@ -// Don't use custom Result type here as we need to return rmcp::ErrorData +use super::error::ClientError; +use super::transport::ClientService; +use super::{ClientOptions, ToolCall, ToolCallOptions, Transport, handler::Handler}; +use futures::future::{BoxFuture, FutureExt, Shared}; +use futures::stream::{self, Stream}; use rmcp::{ - ClientHandler, RoleClient, - handler::client::progress::ProgressDispatcher, - model::{ - ClientCapabilities, ClientConfig, CustomNotification, ElicitRequestParams, ElicitResult, ElicitationAction, - ElicitationCapability, ErrorData, FormElicitationCapability, ProgressNotificationParam, - UrlElicitationCapability, - }, - service::{NotificationContext, RequestContext}, + RoleClient, + model::{CallToolRequestParams, GetPromptRequestParams, GetPromptResult, Prompt, Tool}, + service::{Peer, QuitReason}, }; -use std::result::Result; -use tokio::sync::{mpsc, oneshot}; - -use crate::client::{ElicitationRequest, McpClientEvent, elicitation::with_meta, manager::ToolListChangedRequest}; - +use serde_json::{Map, Value}; +use std::fmt; +use std::pin::pin; +use std::sync::Arc; +use std::time::Duration; +use tokio::task::JoinError; +use tokio::time::timeout; +use tokio_util::sync::{CancellationToken, DropGuard}; + +#[derive(Clone)] pub struct McpClient { - client_info: ClientConfig, - server_name: String, - pub(crate) progress_dispatcher: ProgressDispatcher, - event_sender: mpsc::Sender, - tool_refresh_sender: Option>, - connection_generation: u64, + pub(super) peer: Peer, + pub(super) handler: Arc, + shutdown: Arc, } impl McpClient { - pub fn new(client_info: ClientConfig, server_name: String, event_sender: mpsc::Sender) -> Self { - Self { - client_info, - server_name, - progress_dispatcher: ProgressDispatcher::new(), - event_sender, - tool_refresh_sender: None, - connection_generation: 0, - } + pub async fn connect( + name: impl Into, + transport: Transport, + options: &ClientOptions, + ) -> Result { + transport.connect(Handler::new(name.into(), options), options).await.map(Self::from_service) } - pub(super) fn with_tool_refresh( - mut self, - sender: mpsc::Sender, - connection_generation: u64, - ) -> Self { - self.tool_refresh_sender = Some(sender); - self.connection_generation = connection_generation; - self + pub async fn authorize( + name: impl Into, + transport: Transport, + options: &ClientOptions, + challenge: Option, + ) -> Result { + transport.authorize(Handler::new(name.into(), options), options, challenge).await.map(Self::from_service) } - pub fn server_name(&self) -> &str { - &self.server_name + pub fn name(&self) -> &str { + self.handler.server() } - /// Dispatch an elicitation request through the shared event channel. - /// - /// Used by both the `create_elicitation` handler and the MRTR round loop - /// in `call_tool_mrtr` to ensure the same user-facing flow. - pub async fn dispatch_elicitation(&self, request: ElicitRequestParams) -> ElicitResult { - let (response_tx, response_rx) = oneshot::channel(); - let elicitation_request = - ElicitationRequest { server_name: self.server_name.clone(), request, response_sender: response_tx }; + pub fn description(&self) -> Option { + let info = self.peer.peer_info()?; + non_empty(info.server_info.as_ref()?.description.as_ref()) + } - if self.event_sender.send(McpClientEvent::Elicitation(Box::new(elicitation_request))).await.is_err() { - return cancel_result(); - } - response_rx.await.unwrap_or_else(|_| cancel_result()) + pub fn instructions(&self) -> Option { + non_empty(self.peer.peer_info()?.instructions.as_ref()) } -} -pub fn cancel_result() -> ElicitResult { - ElicitResult::new(ElicitationAction::Cancel) -} + pub async fn list_tools(&self) -> Result, ClientError> { + self.peer.list_all_tools().await.map_err(|source| ClientError::request(self.name(), source)) + } -pub fn client_capabilities() -> ClientCapabilities { - client_capabilities_for(true, true) -} + pub async fn list_prompts(&self) -> Result, ClientError> { + if !self.supports_prompts() { + return Ok(Vec::new()); + } + self.peer.list_all_prompts().await.map_err(|source| ClientError::request(self.name(), source)) + } -pub fn client_capabilities_for(form: bool, url: bool) -> ClientCapabilities { - let mut capabilities = ClientCapabilities::builder().enable_tasks().build(); - if form || url { - let mut elicitation = ElicitationCapability::new(); - elicitation.form = form.then(FormElicitationCapability::default); - elicitation.url = url.then(UrlElicitationCapability::default); - capabilities.elicitation = Some(elicitation); + pub async fn get_prompt(&self, name: &str, args: Map) -> Result { + let mut request = GetPromptRequestParams::new(name); + if !args.is_empty() { + request = request.with_arguments(args); + } + self.peer.get_prompt(request).await.map_err(|source| ClientError::request(self.name(), source)) } - capabilities -} -impl ClientHandler for McpClient { - fn get_info(&self) -> ClientConfig { - self.client_info.clone() + pub fn tool_lists(&self) -> impl Stream, ClientError>> + Send + use<> { + let mut changes = self.handler.tool_list_changes(); + changes.mark_changed(); + stream::unfold((self.clone(), changes), |(client, mut changes)| async move { + tokio::select! { + biased; + () = client.shutdown.closed.clone() => return None, + notified = changes.changed() => notified.ok()?, + } + let tools = client.list_tools().await; + Some((tools, (client, changes))) + }) } - async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext) -> () { - self.progress_dispatcher.handle_notification(params).await; + pub fn call_tool(&self, name: &str, args: Map, options: ToolCallOptions) -> ToolCall { + let params = CallToolRequestParams::new(name.to_string()).with_arguments(args); + ToolCall::new(self.clone(), params, options) } - async fn create_elicitation( - &self, - request: ElicitRequestParams, - context: RequestContext, - ) -> Result { - let meta = (!context.meta.is_empty()).then_some(context.meta); - Ok(self.dispatch_elicitation(with_meta(request, meta)).await) + pub async fn close(&self) { + self.shutdown.cancel.cancel(); + self.shutdown.closed.clone().await; } - async fn on_custom_notification( - &self, - notification: CustomNotification, - _context: NotificationContext, - ) { - if notification.method != "notifications/elicitation/complete" { - return; - } - let params: Option = - notification.params.and_then(|params| serde_json::from_value(params).ok()); - let Some(params) = params else { - tracing::warn!("Ignoring malformed MCP elicitation completion notification"); - return; - }; - let _ = self - .event_sender - .send(McpClientEvent::ElicitationComplete { - server_name: self.server_name.clone(), - elicitation_id: params.elicitation_id, - }) - .await; - } - - async fn on_tool_list_changed(&self, context: NotificationContext) { - let Some(sender) = &self.tool_refresh_sender else { - return; - }; - let request = ToolListChangedRequest::new(self.server_name.clone(), self.connection_generation, context.peer); - if sender.send(request).await.is_err() { - tracing::debug!(server = %self.server_name, "MCP tool refresh receiver closed"); - } + pub(crate) fn close_on_drop(&self) -> DropGuard { + self.shutdown.cancel.clone().drop_guard() } -} -#[derive(serde::Deserialize)] -#[serde(rename_all = "camelCase")] -struct ElicitationCompleteParams { - elicitation_id: String, -} + #[cfg(any(test, feature = "testing"))] + pub fn peer(&self) -> &Peer { + &self.peer + } -#[cfg(test)] -mod tests { - use super::*; - use rmcp::model::{ElicitationSchema, Implementation}; - use std::collections::BTreeMap; + fn supports_prompts(&self) -> bool { + self.peer.peer_info().is_some_and(|info| info.capabilities.prompts.is_some()) + } - fn test_client_info() -> ClientConfig { - ClientConfig::new(client_capabilities(), Implementation::new("test", "0.1.0")) + fn from_service(service: ClientService) -> Self { + let peer = service.peer().clone(); + let handler = Arc::clone(service.service()); + let cancel = CancellationToken::new(); + let closed = tokio::spawn(run_until_closed(service, cancel.clone())).map(|_| ()).boxed().shared(); + let shutdown = Shutdown { _cancel_on_drop: cancel.clone().drop_guard(), cancel, closed }; + Self { peer, handler, shutdown: Arc::new(shutdown) } } +} - fn make_client(event_sender: mpsc::Sender) -> McpClient { - McpClient::new(test_client_info(), "test-server".to_string(), event_sender) +impl fmt::Debug for McpClient { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("McpClient").field("server", &self.name()).finish_non_exhaustive() } +} - fn unwrap_elicitation(event: McpClientEvent) -> ElicitationRequest { - match event { - McpClientEvent::Elicitation(req) => *req, - other => panic!("expected Elicitation, got {other:?}"), +struct Shutdown { + cancel: CancellationToken, + _cancel_on_drop: DropGuard, + closed: Shared>, +} + +const CLOSE_TIMEOUT: Duration = Duration::from_secs(5); + +async fn run_until_closed(service: ClientService, cancel: CancellationToken) { + let server = service.service().server().to_string(); + let stop = service.cancellation_token(); + let mut running = pin!(service.waiting()); + let closed = tokio::select! { + closed = &mut running => closed.map(Some), + () = cancel.cancelled() => { + stop.cancel(); + match timeout(CLOSE_TIMEOUT, running).await { + Ok(closed) => closed.map(Some), + Err(_) => Ok(None), + } } - } + }; + log_close(&server, closed); +} + +fn non_empty(text: Option<&String>) -> Option { + text.filter(|text| !text.is_empty()).cloned() +} - #[tokio::test] - async fn dispatch_elicitation_dropped_sender_returns_cancel() { - let (event_tx, _) = mpsc::channel(1); - let client = make_client(event_tx); - - let request = ElicitRequestParams::FormElicitationParams { - meta: None, - message: "test".to_string(), - requested_schema: ElicitationSchema::new(BTreeMap::new()), - }; - - let result = client.dispatch_elicitation(request).await; - assert_eq!(result.action, ElicitationAction::Cancel, "dropped sender should return Cancel, not Decline"); - assert!(result.content.is_none()); - } - - #[tokio::test] - async fn dispatch_elicitation_dropped_receiver_returns_cancel() { - let (event_tx, mut event_rx) = mpsc::channel(1); - let client = make_client(event_tx); - - let request = ElicitRequestParams::FormElicitationParams { - meta: None, - message: "test".to_string(), - requested_schema: ElicitationSchema::new(BTreeMap::new()), - }; - - let handle = tokio::spawn(async move { - let event = event_rx.recv().await.unwrap(); - let elicitation = unwrap_elicitation(event); - drop(elicitation.response_sender); - }); - - let result = client.dispatch_elicitation(request).await; - handle.await.unwrap(); - - assert_eq!(result.action, ElicitationAction::Cancel, "dropped receiver should return Cancel, not Decline"); - assert!(result.content.is_none()); - } - - #[tokio::test] - async fn dispatch_elicitation_forwards_request_with_server_name() { - let (event_tx, mut event_rx) = mpsc::channel(1); - let client = make_client(event_tx); - - let request = ElicitRequestParams::UrlElicitationParams { - meta: None, - message: "Auth".to_string(), - url: "https://example.com/auth".to_string(), - elicitation_id: "el-123".to_string(), - }; - - let handle = tokio::spawn(async move { - let event = event_rx.recv().await.unwrap(); - let elicitation = unwrap_elicitation(event); - assert_eq!(elicitation.server_name, "test-server"); - let _ = elicitation.response_sender.send(ElicitResult::new(ElicitationAction::Accept)); - }); - - let result = client.dispatch_elicitation(request).await; - handle.await.unwrap(); - assert_eq!(result.action, ElicitationAction::Accept); - } - - #[test] - fn capabilities_include_form_url_and_tasks() { - let info = test_client_info(); - let caps = &info.capabilities; - let elicitation = caps.elicitation.as_ref().expect("elicitation capability should be set"); - assert!(elicitation.form.is_some(), "form capability should be advertised"); - assert!(elicitation.url.is_some(), "url capability should be advertised"); - assert!( - caps.extensions.as_ref().is_some_and(|extensions| extensions.contains_key("io.modelcontextprotocol/tasks")) - ); +fn log_close(server: &str, result: Result, JoinError>) { + match result { + Ok(Some(reason)) => tracing::debug!(server, ?reason, "MCP server connection closed"), + Ok(None) => tracing::warn!(server, "MCP server connection did not close within {CLOSE_TIMEOUT:?}"), + Err(error) => tracing::warn!(server, %error, "MCP server connection task failed while closing"), } } diff --git a/crates/mcp-utils/src/client/mcp_snapshot.rs b/crates/mcp-utils/src/client/mcp_snapshot.rs deleted file mode 100644 index df483f70b..000000000 --- a/crates/mcp-utils/src/client/mcp_snapshot.rs +++ /dev/null @@ -1,100 +0,0 @@ -use super::{ - McpClient, McpError, ToolCatalog, ToolRoute, - naming::{create_namespaced_tool_name, split_on_server_name}, -}; -use llm::ToolDefinition; -use rmcp::{RoleClient, model::CallToolRequestParams, service::RunningService}; -use serde_json::{Map, Value}; -use std::{collections::HashMap, fmt, sync::Arc}; - -#[derive(Clone, Default)] -pub struct McpSnapshot { - catalog: Arc, - clients: Arc>>>, -} - -impl fmt::Debug for McpSnapshot { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("McpSnapshot") - .field("catalog", &self.catalog) - .field("connected_servers", &self.clients.keys().collect::>()) - .finish() - } -} - -impl McpSnapshot { - pub fn new( - catalog: Arc, - clients: Arc>>>, - ) -> Self { - Self { catalog, clients } - } - - pub fn catalog(&self) -> &Arc { - &self.catalog - } - - pub fn tool_definitions(&self) -> Vec { - self.catalog.tools().model_visible.into_iter().map(|tool| tool.definition().clone()).collect() - } - - pub fn model_instructions(&self) -> std::collections::BTreeMap { - self.catalog.model_instructions() - } - - pub fn server_statuses(&self) -> Vec { - self.catalog.server_statuses() - } - - pub fn resolve( - &self, - route: ToolRoute, - arguments: Map, - ) -> super::Result<(Arc>, CallToolRequestParams)> { - if let ToolRoute::ModelVisible { namespaced_name } = &route { - split_on_server_name(namespaced_name) - .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_name.clone()))?; - } - if !self.catalog.route_permitted(&route) { - let (tool_name, namespaced_name) = match &route { - ToolRoute::ModelVisible { namespaced_name } => (namespaced_name.clone(), namespaced_name.clone()), - ToolRoute::Deferred { server, tool } => (tool.clone(), create_namespaced_tool_name(server, tool)), - }; - if matches!(route, ToolRoute::Deferred { .. }) - && self.catalog.route_permitted(&ToolRoute::ModelVisible { namespaced_name: namespaced_name.clone() }) - { - return Err(McpError::DirectToolRequiresDirectRoute { tool_name, direct_name: namespaced_name }); - } - return Err(McpError::ToolNotFound(namespaced_name)); - } - let (server, tool) = match route { - ToolRoute::ModelVisible { namespaced_name } => { - let (server, tool) = - split_on_server_name(&namespaced_name).expect("model-visible route was validated above"); - (server.to_string(), tool.to_string()) - } - ToolRoute::Deferred { server, tool } => (server, tool), - }; - let client = self.clients.get(&server).cloned().ok_or_else(|| McpError::ServerNotFound(server.clone()))?; - Ok((client, CallToolRequestParams::new(tool).with_arguments(arguments))) - } - - pub fn clients_with_prompts(&self) -> Vec<(String, Arc>)> { - self.clients - .iter() - .filter(|(_, client)| client.peer_info().is_some_and(|info| info.capabilities.prompts.is_some())) - .map(|(name, client)| (name.clone(), Arc::clone(client))) - .collect() - } - - pub fn client_for_prompt( - &self, - namespaced_name: &str, - ) -> super::Result<(String, Arc>)> { - let (server, prompt) = split_on_server_name(namespaced_name) - .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_name.to_string()))?; - let client = self.clients.get(server).cloned().ok_or_else(|| McpError::ServerNotFound(server.to_string()))?; - Ok((prompt.to_string(), client)) - } -} diff --git a/crates/mcp-utils/src/client/mod.rs b/crates/mcp-utils/src/client/mod.rs index 3bec98899..6c6a68972 100644 --- a/crates/mcp-utils/src/client/mod.rs +++ b/crates/mcp-utils/src/client/mod.rs @@ -1,43 +1,24 @@ -pub mod config; -pub mod error; -pub mod manager; -pub mod oauth_handler; +//! [`McpClient`]: one live connection to one MCP server. -mod call_tool; -mod connection; -mod connection_attempt_manager; mod elicitation; +mod error; +mod handler; mod mcp_client; -mod mcp_snapshot; -mod mrtr; -mod naming; +mod oauth; +mod options; mod task; -mod tool_catalog; -mod tool_filter; +mod tool_call; +mod transport; -pub use call_tool::{CallToolError, CallToolOptions, ToolCallEvent, call_tool}; -pub use config::{ - AETHER_OAUTH_CALLBACK_PORT, AETHER_OAUTH_CLIENT_METADATA_URL, DeferredToolRules, InMemoryServerConfig, - InMemoryServerSpec, InMemoryType, McpConfig, McpHttpConfig, McpOAuthConfig, McpServer, McpServerConfig, - McpTransport, ParseError, RemoteServerConfig, RemoteType, ResolvedOAuth, StdioServerConfig, StdioType, - ToolExposure, loopback_redirect_uri, -}; -pub use connection::{McpConnectAttempt, McpConnectOutcome, McpServerConnection}; -pub use connection_attempt_manager::McpConnectionAttemptManager; -pub use error::{McpError, Result}; -pub use manager::{ - ElicitationRequest, McpClientEvent, McpConnectionDetails, McpManager, OAuthHandlerContext, OAuthHandlerFactory, - RuntimeMcpServer, RuntimeMcpTransport, -}; -pub use mcp_client::{McpClient, cancel_result, client_capabilities, client_capabilities_for}; -pub use mcp_snapshot::McpSnapshot; -pub use mrtr::AbortReason; -pub use naming::{SERVERNAME_DELIMITER, split_on_server_name}; -pub use oauth_handler::ElicitingOAuthHandler; +#[cfg(any(test, feature = "testing"))] +pub(crate) use elicitation::cancelled; +pub use elicitation::{Elicitation, ElicitationRequest, ElicitationSink}; +pub use error::ClientError; +pub use mcp_client::McpClient; +pub use oauth::OAuthSettings; +pub use options::ClientOptions; +pub use reqwest::header::HeaderMap; pub use task::TaskErrorReason; pub use tokio_util::sync::CancellationToken; -pub use tool_catalog::{ - CatalogTool, CatalogTools, PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME, ServerCatalogEntry, ServerDescription, - ToolCatalog, ToolExposureKind, ToolRoute, -}; -pub use tool_filter::{ToolAnnotationMatcher, ToolFilter, ToolMatcher}; +pub use tool_call::{ToolCall, ToolCallError, ToolCallEvent, ToolCallOptions}; +pub use transport::Transport; diff --git a/crates/mcp-utils/src/client/mrtr.rs b/crates/mcp-utils/src/client/mrtr.rs deleted file mode 100644 index 50dc9ab57..000000000 --- a/crates/mcp-utils/src/client/mrtr.rs +++ /dev/null @@ -1,285 +0,0 @@ -use rmcp::model::{DEFAULT_MRTR_MAX_ROUNDS, ElicitResult, ElicitationAction, InputRequests, InputRequiredResult}; -use std::time::Duration; - -#[derive(Debug)] -pub struct MrtrState { - timeout: Duration, - input_request_rounds: usize, - next_backoff: Duration, - state_only_waited: Duration, - user_cancelled: bool, -} - -#[derive(Debug, PartialEq)] -pub enum MrtrAction { - /// Sleep for the backoff, then retry with the echoed request state. - Poll { backoff: Duration, request_state: String }, - /// Dispatch each input request to the user, recording every response via - /// [`MrtrState::record_response`], then retry with the responses. - Elicit { input_requests: InputRequests, request_state: Option }, - /// Fail the tool call. - Abort(AbortReason), -} - -#[derive(Debug, PartialEq, Eq)] -pub enum AbortReason { - EmptyInputRequired, - PollingBudgetExhausted, - RePromptAfterCancel, - InputRoundsExceeded, -} - -impl MrtrState { - pub fn new(timeout: Duration) -> Self { - Self { - timeout, - input_request_rounds: 0, - next_backoff: BASE_BACKOFF, - state_only_waited: Duration::ZERO, - user_cancelled: false, - } - } - - pub fn tick(&mut self, input_required: InputRequiredResult) -> MrtrAction { - let input_requests = input_required.input_requests.filter(|requests| !requests.is_empty()); - match (input_requests, input_required.request_state) { - (None, None) => MrtrAction::Abort(AbortReason::EmptyInputRequired), - (None, Some(request_state)) => { - let backoff = self.next_backoff; - if self.state_only_waited + backoff > self.timeout { - MrtrAction::Abort(AbortReason::PollingBudgetExhausted) - } else { - self.state_only_waited += backoff; - self.next_backoff = (backoff * 2).min(MAX_BACKOFF); - MrtrAction::Poll { backoff, request_state } - } - } - (Some(input_requests), request_state) => { - if self.user_cancelled { - MrtrAction::Abort(AbortReason::RePromptAfterCancel) - } else if self.input_request_rounds == DEFAULT_MRTR_MAX_ROUNDS { - MrtrAction::Abort(AbortReason::InputRoundsExceeded) - } else { - self.input_request_rounds += 1; - self.next_backoff = BASE_BACKOFF; - self.state_only_waited = Duration::ZERO; - MrtrAction::Elicit { input_requests, request_state } - } - } - } - } - - /// Record an elicitation response so a user cancellation makes the next input-requesting - /// [`MrtrState::tick`] abort instead of re-prompting. - pub fn record_response(&mut self, response: &ElicitResult) { - self.user_cancelled |= response.action == ElicitationAction::Cancel; - } -} - -impl AbortReason { - pub fn message(&self, server_name: &str, timeout: Duration) -> String { - match self { - Self::EmptyInputRequired => { - format!("Server '{server_name}' requested input without any input requests or request state") - } - Self::PollingBudgetExhausted => { - format!("Server '{server_name}' did not complete within {}s of state-only polling", timeout.as_secs()) - } - Self::RePromptAfterCancel => { - format!("Input requested by server '{server_name}' was cancelled by the user") - } - Self::InputRoundsExceeded => { - format!("Server '{server_name}' did not complete within {DEFAULT_MRTR_MAX_ROUNDS} MRTR input rounds") - } - } - } -} - -const BASE_BACKOFF: Duration = Duration::from_millis(50); -const MAX_BACKOFF: Duration = Duration::from_millis(1600); - -#[cfg(test)] -mod tests { - use super::*; - use rmcp::model::{ElicitRequest, ElicitRequestParams, InputRequest}; - - const TIMEOUT: Duration = Duration::from_secs(1); - - #[test] - fn empty_input_required_aborts() { - let mut test = mrtr(); - let trace = test.run(|_| input_required().build()); - assert!(matches!(trace.last(), Some(MrtrAction::Abort(AbortReason::EmptyInputRequired)))); - } - - #[test] - fn state_only_polls_with_growing_backoff_until_the_budget_is_spent() { - let timeout = Duration::from_secs(4); - let mut test = mrtr().with_timeout(timeout); - let trace = test.run(|_| input_required().with_request_state("s").build()); - - assert_eq!( - trace.as_slice(), - &[ - MrtrAction::Poll { backoff: Duration::from_millis(50), request_state: "s".to_string() }, - MrtrAction::Poll { backoff: Duration::from_millis(100), request_state: "s".to_string() }, - MrtrAction::Poll { backoff: Duration::from_millis(200), request_state: "s".to_string() }, - MrtrAction::Poll { backoff: Duration::from_millis(400), request_state: "s".to_string() }, - MrtrAction::Poll { backoff: Duration::from_millis(800), request_state: "s".to_string() }, - MrtrAction::Poll { backoff: Duration::from_millis(1600), request_state: "s".to_string() }, - MrtrAction::Abort(AbortReason::PollingBudgetExhausted), - ] - ); - } - - #[test] - fn an_input_round_refreshes_the_polling_budget() { - let mut test = mrtr(); - while matches!(test.rounds.tick(input_required().with_request_state("s").build()), MrtrAction::Poll { .. }) {} - test.elicit_round(); - - let decision = test.rounds.tick(input_required().with_request_state("s").build()); - assert!( - matches!(decision, MrtrAction::Poll { backoff, .. } if backoff == BASE_BACKOFF), - "budget and backoff should reset after an input round, got {decision:?}" - ); - } - - #[test] - fn input_rounds_abort_at_the_cap() { - let mut test = mrtr(); - let actions = test.run(|_| input_required().with_form().build()); - assert!( - actions.as_slice()[..DEFAULT_MRTR_MAX_ROUNDS] - .iter() - .all(|decision| matches!(decision, MrtrAction::Elicit { .. })) - ); - assert!(matches!(actions.last(), Some(MrtrAction::Abort(AbortReason::InputRoundsExceeded)))); - } - - #[test] - fn a_cancelled_round_aborts_the_next_prompt_but_not_state_only_polling() { - let mut test = mrtr().answering_with(ElicitationAction::Cancel); - let actions = test.run(|_| input_required().with_form().build()); - assert!(matches!( - actions.as_slice(), - [MrtrAction::Elicit { .. }, MrtrAction::Abort(AbortReason::RePromptAfterCancel)] - )); - - let poll = test.rounds.tick(input_required().with_request_state("s").build()); - assert!(matches!(poll, MrtrAction::Poll { .. }), "the server may still finish up, got {poll:?}"); - } - - #[test] - fn an_accepted_round_allows_the_next_prompt() { - let mut test = mrtr().answering_with(ElicitationAction::Accept); - let elicitation_count = std::cell::Cell::new(0); - let actions = test.run_until( - |_| { - elicitation_count.set(elicitation_count.get() + 1); - input_required().with_form().build() - }, - |decision| matches!(decision, MrtrAction::Elicit { .. }) && elicitation_count.get() == 2, - ); - - assert!(matches!(actions.as_slice(), [MrtrAction::Elicit { .. }, MrtrAction::Elicit { .. }])); - } - - struct MrtrTest { - rounds: MrtrState, - response: ElicitResult, - } - - fn mrtr() -> MrtrTest { - MrtrTest::default() - } - - impl Default for MrtrTest { - fn default() -> Self { - Self { rounds: MrtrState::new(TIMEOUT), response: ElicitResult::new(ElicitationAction::Accept) } - } - } - - impl MrtrTest { - fn with_timeout(mut self, timeout: Duration) -> Self { - self.rounds = MrtrState::new(timeout); - self - } - - fn answering_with(mut self, action: ElicitationAction) -> Self { - self.response = ElicitResult::new(action); - self - } - - fn run_until(&mut self, mut next: T, mut terminal: U) -> Vec - where - T: FnMut(Option<&MrtrAction>) -> InputRequiredResult, - U: FnMut(&MrtrAction) -> bool, - { - let mut actions = Vec::new(); - loop { - let action = self.rounds.tick(next(actions.last())); - if matches!(action, MrtrAction::Elicit { .. }) { - self.rounds.record_response(&self.response); - } - let is_terminal = terminal(&action); - actions.push(action); - if is_terminal { - return actions; - } - } - } - - fn run(&mut self, next: T) -> Vec - where - T: FnMut(Option<&MrtrAction>) -> InputRequiredResult, - { - self.run_until(next, |decision| matches!(decision, MrtrAction::Abort(_))) - } - - fn elicit_round(&mut self) { - let actions = self.run_until( - |_| input_required().with_form().build(), - |decision| matches!(decision, MrtrAction::Elicit { .. }), - ); - assert!(matches!(actions.last(), Some(MrtrAction::Elicit { .. }))); - } - } - - fn input_required() -> InputRequiredBuilder { - InputRequiredBuilder::default() - } - - #[derive(Default)] - struct InputRequiredBuilder { - input_requests: Option, - request_state: Option, - } - - impl InputRequiredBuilder { - fn with_form(mut self) -> Self { - let params = ElicitRequestParams::FormElicitationParams { - meta: None, - message: "m".to_string(), - requested_schema: serde_json::from_value(serde_json::json!({ - "type": "object", - "properties": {} - })) - .unwrap(), - }; - let mut requests = InputRequests::new(); - requests.insert("k".to_string(), InputRequest::Elicitation(ElicitRequest::new(params))); - self.input_requests = Some(requests); - self - } - - fn with_request_state(mut self, state: &str) -> Self { - self.request_state = Some(state.to_string()); - self - } - - fn build(self) -> InputRequiredResult { - InputRequiredResult::new(self.input_requests, self.request_state) - } - } -} diff --git a/crates/mcp-utils/src/client/naming.rs b/crates/mcp-utils/src/client/naming.rs deleted file mode 100644 index abb18f721..000000000 --- a/crates/mcp-utils/src/client/naming.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub const SERVERNAME_DELIMITER: &str = "__"; - -pub(crate) fn create_namespaced_tool_name(server_name: &str, tool_name: &str) -> String { - format!("{server_name}{SERVERNAME_DELIMITER}{tool_name}") -} - -pub fn split_on_server_name(namespaced_name: &str) -> Option<(&str, &str)> { - namespaced_name.split_once(SERVERNAME_DELIMITER) -} diff --git a/crates/mcp-utils/src/client/oauth.rs b/crates/mcp-utils/src/client/oauth.rs new file mode 100644 index 000000000..b3dc00b56 --- /dev/null +++ b/crates/mcp-utils/src/client/oauth.rs @@ -0,0 +1,137 @@ +use super::error::{ClientError, Result}; +use super::{ClientOptions, Elicitation, elicitation::ElicitationRequest}; +use aether_auth::{ + OAuthClientRegistration, OAuthCredentialStorage, OAuthError, OAuthFlowOptions, OAuthHandler, accept_oauth_callback, + create_auth_manager_from_store, perform_oauth_flow, +}; +use futures::future::BoxFuture; +use rmcp::model::{ElicitRequestParams, ElicitationAction}; +use rmcp::transport::auth::AuthClient; +use std::sync::Arc; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::sync::mpsc; +use tokio::time::timeout; + +#[derive(Debug, Clone, Default)] +pub struct OAuthSettings { + pub registration: OAuthClientRegistration, + pub callback_port: u16, +} + +pub(super) async fn restore( + server: &str, + url: &str, + oauth: &OAuthSettings, + store: &Arc, +) -> Option> { + let restored = create_auth_manager_from_store( + server, + url, + oauth.registration.pre_registered_client_id(), + &loopback_redirect_uri(oauth.callback_port), + Arc::clone(store), + ) + .await; + match restored { + Ok(manager) => manager.map(|manager| AuthClient::new(reqwest::Client::default(), manager)), + Err(error) => { + tracing::warn!( + server, + %error, + "Failed to initialize auth manager from stored credentials, proceeding without auth" + ); + None + } + } +} + +pub(super) async fn authorize( + server: &str, + url: &str, + oauth: OAuthSettings, + options: &ClientOptions, + challenge: Option, +) -> Result> { + let events = options.oauth_prompts().ok_or_else(|| ClientError::OAuthUnavailable { server: server.to_string() })?; + let auth_error = |source| ClientError::Auth { server: server.to_string(), source }; + let handler = ElicitingOAuthHandler::bind(server, oauth.callback_port, events.clone()) + .map_err(|error| auth_error(OAuthError::from(error)))?; + let flow = perform_oauth_flow( + server, + url, + &handler, + OAuthFlowOptions { client_registration: oauth.registration, challenge }, + options.oauth_store.clone(), + ); + timeout(AUTH_TIMEOUT, flow) + .await + .map_err(|_| ClientError::AuthTimedOut { server: server.to_string(), timeout: AUTH_TIMEOUT })? + .map_err(auth_error) +} + +const AUTH_TIMEOUT: Duration = Duration::from_mins(3); +const OAUTH_ELICITATION_ID: &str = "oauth"; + +struct ElicitingOAuthHandler { + listener: TcpListener, + redirect_uri: String, + server: String, + events: mpsc::Sender, +} + +impl ElicitingOAuthHandler { + fn bind(server: &str, port: u16, events: mpsc::Sender) -> std::io::Result { + let listener = std::net::TcpListener::bind(("127.0.0.1", port))?; + let port = listener.local_addr()?.port(); + listener.set_nonblocking(true)?; + Ok(Self { + listener: TcpListener::from_std(listener)?, + redirect_uri: loopback_redirect_uri(port), + server: server.to_string(), + events, + }) + } +} + +impl OAuthHandler for ElicitingOAuthHandler { + fn redirect_uri(&self) -> &str { + &self.redirect_uri + } + + fn authorize(&self, auth_url: &str) -> BoxFuture<'_, std::result::Result> { + let request = ElicitRequestParams::UrlElicitationParams { + meta: None, + message: "Open this URL to authorize MCP server access.".to_string(), + url: auth_url.to_string(), + elicitation_id: OAUTH_ELICITATION_ID.to_string(), + }; + Box::pin(async move { + let (elicitation, response) = ElicitationRequest::new(self.server.clone(), request); + self.events + .send(Elicitation::Request(Box::new(elicitation))) + .await + .map_err(|_| OAuthError::Rmcp("OAuth prompt channel closed".to_string()))?; + + let result = tokio::select! { + callback = accept_oauth_callback(&self.listener) => callback, + response = response => match response { + Ok(result) if result.action == ElicitationAction::Accept => { + accept_oauth_callback(&self.listener).await + } + Ok(_) | Err(_) => Err(OAuthError::UserCancelled), + }, + }; + if !matches!(result, Err(OAuthError::UserCancelled)) { + let completion = + Elicitation::Complete { server: self.server.clone(), id: OAUTH_ELICITATION_ID.to_string() }; + let _ = self.events.send(completion).await; + } + result + }) + } +} + +fn loopback_redirect_uri(port: u16) -> String { + format!("http://localhost:{port}/") +} diff --git a/crates/mcp-utils/src/client/oauth_handler.rs b/crates/mcp-utils/src/client/oauth_handler.rs deleted file mode 100644 index b8786e670..000000000 --- a/crates/mcp-utils/src/client/oauth_handler.rs +++ /dev/null @@ -1,152 +0,0 @@ -use crate::client::config::loopback_redirect_uri; -use crate::client::manager::{ElicitationRequest, McpClientEvent, OAuthHandlerContext}; -use aether_auth::{OAuthError, OAuthHandler, accept_oauth_callback}; -use futures::future::BoxFuture; -use rmcp::model::{ElicitRequestParams, ElicitationAction}; -use std::num::NonZeroU16; -use tokio::net::TcpListener; -use tokio::sync::{mpsc, oneshot}; - -const AETHER_OAUTH_ELICITATION_ID: &str = "aether-oauth"; - -/// `OAuthHandler` that dispatches the OAuth authorization URL to the host. -pub struct ElicitingOAuthHandler { - listener: TcpListener, - redirect_uri: String, - server_name: String, - event_sender: mpsc::Sender, -} - -impl ElicitingOAuthHandler { - pub fn new(ctx: OAuthHandlerContext) -> Result { - let port = ctx.callback_port.map_or(0, NonZeroU16::get); - let std_listener = std::net::TcpListener::bind(("127.0.0.1", port))?; - let port = std_listener.local_addr()?.port(); - std_listener.set_nonblocking(true)?; - Ok(Self { - listener: TcpListener::from_std(std_listener)?, - redirect_uri: loopback_redirect_uri(port), - server_name: ctx.server_name, - event_sender: ctx.tx, - }) - } -} - -impl OAuthHandler for ElicitingOAuthHandler { - fn redirect_uri(&self) -> &str { - &self.redirect_uri - } - - fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { - let auth_url = auth_url.to_string(); - Box::pin(async move { - let (response_sender, response_rx) = oneshot::channel(); - self.event_sender - .send(McpClientEvent::Elicitation(Box::new(ElicitationRequest { - server_name: self.server_name.clone(), - request: ElicitRequestParams::UrlElicitationParams { - meta: None, - message: "Open this URL to authorize MCP server access.".to_string(), - url: auth_url, - elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(), - }, - response_sender, - }))) - .await - .map_err(|_| OAuthError::Rmcp("OAuth prompt channel closed".to_string()))?; - - let result = tokio::select! { - callback = accept_oauth_callback(&self.listener) => callback, - response = response_rx => match response { - Ok(result) if matches!(result.action, ElicitationAction::Decline | ElicitationAction::Cancel) => { - Err(OAuthError::UserCancelled) - } - Ok(_) | Err(_) => accept_oauth_callback(&self.listener).await, - }, - }; - if !matches!(result, Err(OAuthError::UserCancelled)) { - let _ = self - .event_sender - .send(McpClientEvent::ElicitationComplete { - server_name: self.server_name.clone(), - elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(), - }) - .await; - } - result - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rmcp::model::ElicitResult; - use std::sync::Arc; - use tokio::{io::AsyncWriteExt, task::yield_now}; - - #[tokio::test] - async fn accepting_browser_prompt_keeps_waiting_for_callback() { - let (tx, mut rx) = mpsc::channel(1); - let handler = Arc::new( - ElicitingOAuthHandler::new(OAuthHandlerContext { - server_name: "slack".to_string(), - callback_port: None, - tx, - }) - .unwrap(), - ); - let port = handler - .redirect_uri() - .strip_prefix("http://localhost:") - .and_then(|value| value.strip_suffix('/')) - .unwrap() - .parse::() - .unwrap(); - let authorize = { - let handler = Arc::clone(&handler); - tokio::spawn(async move { handler.authorize("https://example.com/oauth").await }) - }; - let McpClientEvent::Elicitation(request) = rx.recv().await.unwrap() else { - panic!("expected OAuth elicitation"); - }; - request.response_sender.send(ElicitResult::new(ElicitationAction::Accept)).unwrap(); - yield_now().await; - assert!(!authorize.is_finished()); - - let mut callback = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap(); - callback.write_all(b"GET /?code=test-code&state=test-state HTTP/1.1\r\nHost: localhost\r\n\r\n").await.unwrap(); - - assert!(authorize.await.unwrap().unwrap().contains("code=test-code&state=test-state")); - } - - #[tokio::test] - async fn configured_callback_port_uses_registered_redirect() { - let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let port = probe.local_addr().unwrap().port(); - drop(probe); - let (tx, _) = mpsc::channel(1); - let handler = ElicitingOAuthHandler::new(OAuthHandlerContext { - server_name: "slack".to_string(), - callback_port: NonZeroU16::new(port), - tx, - }) - .unwrap(); - assert_eq!(handler.redirect_uri(), format!("http://localhost:{port}/")); - } - - #[tokio::test] - async fn configured_callback_port_fails_when_in_use() { - let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - let port = listener.local_addr().unwrap().port(); - let (tx, _) = mpsc::channel(1); - let error = ElicitingOAuthHandler::new(OAuthHandlerContext { - server_name: "slack".to_string(), - callback_port: NonZeroU16::new(port), - tx, - }) - .err() - .unwrap(); - assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse); - } -} diff --git a/crates/mcp-utils/src/client/options.rs b/crates/mcp-utils/src/client/options.rs new file mode 100644 index 000000000..decea0377 --- /dev/null +++ b/crates/mcp-utils/src/client/options.rs @@ -0,0 +1,69 @@ +use super::{Elicitation, ElicitationSink}; +use aether_auth::OAuthCredentialStorage; +use rmcp::model::{ClientCapabilities, ClientConfig, Implementation}; +use std::sync::Arc; +use tokio::sync::mpsc; + +#[derive(Clone)] +pub struct ClientOptions { + pub(super) implementation: Implementation, + pub(super) elicitation: Option, + pub(super) oauth_store: Option>, + tasks: bool, +} + +impl Default for ClientOptions { + fn default() -> Self { + Self { + implementation: Implementation::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION")), + elicitation: None, + oauth_store: None, + tasks: true, + } + } +} + +impl ClientOptions { + pub fn implementation(mut self, implementation: Implementation) -> Self { + self.implementation = implementation; + self + } + + pub fn elicitation(mut self, sink: ElicitationSink) -> Self { + self.elicitation = Some(sink); + self + } + + pub fn oauth_store(mut self, store: Arc) -> Self { + self.oauth_store = Some(store); + self + } + + #[cfg(any(test, feature = "testing"))] + pub fn without_tasks(mut self) -> Self { + self.tasks = false; + self + } + + pub(super) fn client_config(&self) -> ClientConfig { + let mut capabilities = if self.tasks { + ClientCapabilities::builder().enable_tasks().build() + } else { + ClientCapabilities::default() + }; + capabilities.elicitation = self + .elicitation + .as_ref() + .map(|sink| sink.capability.clone()) + .filter(|capability| capability.form.is_some() || capability.url.is_some()); + ClientConfig::new(capabilities, self.implementation.clone()) + } + + pub(super) fn elicitations(&self) -> Option<&mpsc::Sender> { + self.elicitation.as_ref().map(|sink| &sink.events) + } + + pub(super) fn oauth_prompts(&self) -> Option<&mpsc::Sender> { + self.elicitation.as_ref().filter(|sink| sink.capability.url.is_some()).map(|sink| &sink.events) + } +} diff --git a/crates/mcp-utils/src/client/task.rs b/crates/mcp-utils/src/client/task.rs index 91a1c73a2..cfe7a203b 100644 --- a/crates/mcp-utils/src/client/task.rs +++ b/crates/mcp-utils/src/client/task.rs @@ -1,23 +1,17 @@ -use crate::client::McpClient; -use crate::client::call_tool::{CallToolError, ToolCallEvent}; -use crate::client::elicitation::{ElicitInputsError, elicit_inputs}; +use super::McpClient; +use super::tool_call::{ToolCallError, ToolCallEvent, with_progress}; use async_stream::stream; -use futures::future::{Either, select}; -use futures::{Stream, StreamExt, pin_mut}; -use rmcp::RoleClient; +use futures::future::Either; +use futures::{Stream, StreamExt}; use rmcp::model::{ CancelTaskParams, CreateTaskResult, GetTaskParams, InputRequests, ProgressNotificationParam, Task, TaskPayload, TaskStatus, UpdateTaskParams, }; -use rmcp::service::{RunningService, ServiceError}; +use rmcp::service::ServiceError; use std::collections::HashSet; -use std::future::Future; -use std::pin::pin; use std::time::Duration; use thiserror::Error; -use tokio::time::error::Elapsed; -use tokio::time::{Instant, sleep, timeout, timeout_at}; -use tokio_util::sync::CancellationToken; +use tokio::time::{sleep, timeout}; #[derive(Debug, Error)] pub enum TaskErrorReason { @@ -27,8 +21,6 @@ pub enum TaskErrorReason { Update(#[source] ServiceError), #[error("expired before completion")] Expired, - #[error("exceeded the {timeout:?} execution deadline")] - TimedOut { timeout: Duration }, #[error("repeated input requests that were already answered")] RepeatedInput, #[error("failed: {error}")] @@ -37,251 +29,122 @@ pub enum TaskErrorReason { Cancelled, #[error("returned a malformed result: {0}")] MalformedResult(#[source] serde_json::Error), - #[error("requested an input kind this client does not support")] - UnsupportedInput, - #[error("produced an elicitation response that could not be serialized: {0}")] - Serialize(#[source] serde_json::Error), #[error("returned a task payload this client does not support (status {status:?})")] UnsupportedPayload { status: TaskStatus }, } -pub(crate) struct TaskDriver<'a> { - server_name: &'a str, - client: &'a RunningService, - timeout: Duration, - cancellation_token: CancellationToken, - default_poll_interval: Duration, +pub(super) fn stream_task_events<'a>( + client: &'a McpClient, + created: CreateTaskResult, + progress: impl Stream + 'a, +) -> impl Stream + 'a { + let polled = with_progress(progress, poll_task(client, created.task.clone())).map(|event| match event { + Either::Left(progress) => ToolCallEvent::Progress(progress), + Either::Right(event) => event, + }); + futures::stream::once(std::future::ready(ToolCallEvent::TaskCreated(created))).chain(polled) } -impl<'a> TaskDriver<'a> { - pub(crate) fn new( - server_name: &'a str, - client: &'a RunningService, - timeout: Duration, - cancellation_token: CancellationToken, - ) -> Self { - Self { client, server_name, timeout, cancellation_token, default_poll_interval: Duration::from_secs(1) } +pub(super) async fn cancel_server_task(client: &McpClient, task_id: &str) { + let server_name = client.name(); + match timeout(Duration::from_secs(1), client.peer.cancel_task(CancelTaskParams::new(task_id))).await { + Ok(Ok(())) => {} + Ok(Err(error)) => { + tracing::warn!(server = %server_name, %task_id, "Failed to cancel abandoned MCP task: {error}"); + } + Err(_) => tracing::warn!(server = %server_name, %task_id, "Timed out cancelling abandoned MCP task"), } +} - pub(crate) fn stream + Send + 'a>( - self, - created: CreateTaskResult, - progress_events: T, - ) -> impl Stream + 'a { - stream! { - yield ToolCallEvent::TaskCreated(created.clone()); - let task_events = self.stream_task_events(created.task); - pin_mut!(task_events); - pin_mut!(progress_events); +fn poll_task(client: &McpClient, mut task: Task) -> impl Stream + '_ { + stream! { + let mut answered = HashSet::new(); - loop { - tokio::select! { - progress_event = progress_events.next() => { - let Some(progress_event) = progress_event else { - while let Some(event) = task_events.next().await { - yield event; - } - return; - }; - yield ToolCallEvent::Progress(progress_event); - } - event = task_events.next() => { - let Some(event) = event else { - return; - }; - yield event; - } - } + loop { + if is_task_expired(&task) { + yield failed(&task.task_id, TaskErrorReason::Expired); + return; } - } - } - - fn stream_task_events(self, mut task: Task) -> impl Stream + 'a { - stream! { - let bounds = TaskBounds::new(self.timeout, self.cancellation_token.clone()); - let mut answered_input_keys = HashSet::new(); - loop { - if is_task_expired(&task) { - yield self.fail(task, TaskErrorReason::Expired); + let detailed_task = match client.peer.get_task(GetTaskParams::new(task.task_id.clone())).await { + Ok(result) => result.task, + Err(error) => { + yield abandoned(client, &task.task_id, TaskErrorReason::Get(error)).await; return; } + }; - let detailed_task = match bounds - .run(self.client.get_task(GetTaskParams::new(task.task_id.clone()))) - .await - { - Ok(Ok(result)) => result.task, - Ok(Err(source)) => { - yield self.cancel(task, TaskErrorReason::Get(source)).await; - return; - } - Err(interrupt) => { - yield self.interrupted(task, interrupt).await; - return; - } - }; - - task = detailed_task.task; - if !task.status.is_terminal() { - yield ToolCallEvent::TaskStatus(task.clone()); - } + task = detailed_task.task; + if !task.status.is_terminal() { + yield ToolCallEvent::TaskStatus(task.clone()); + } - match detailed_task.payload { - TaskPayload::Working => {} - TaskPayload::InputRequired { input_requests } => { - match bounds.run(self.elicit_inputs(input_requests, &mut answered_input_keys, &task.task_id)).await { - Ok(Ok(())) => {} - Ok(Err(reason)) => { - yield self.cancel(task, reason).await; - return; - } - Err(interrupt) => { - yield self.interrupted(task, interrupt).await; - return; - } - } - } - TaskPayload::Completed { result } => { - let result = serde_json::from_value(serde_json::Value::Object(result)) - .map_err(|source| self.error(&task.task_id, TaskErrorReason::MalformedResult(source))); - yield ToolCallEvent::TaskComplete { task, result }; - return; - } - TaskPayload::Failed { error } => { - yield self.fail(task, TaskErrorReason::Failed { error: serde_json::Value::Object(error) }); - return; - } - TaskPayload::Cancelled => { - yield self.fail(task, TaskErrorReason::Cancelled); - return; - } - _ => { - let status = task.status; - yield self.cancel(task, TaskErrorReason::UnsupportedPayload { status }).await; + match detailed_task.payload { + TaskPayload::Working => {} + TaskPayload::InputRequired { input_requests } => { + if let Err(error) = answer_inputs(client, &task.task_id, input_requests, &mut answered).await { + cancel_server_task(client, &task.task_id).await; + yield error.into(); return; } } - - let duration = task.poll_interval_ms.map_or(self.default_poll_interval, Duration::from_millis); - if let Err(interrupt) = bounds.run(sleep(duration)).await { - yield self.interrupted(task, interrupt).await; + TaskPayload::Completed { result } => { + let result = serde_json::from_value(serde_json::Value::Object(result)) + .map_err(|source| task_error(&task.task_id, TaskErrorReason::MalformedResult(source))); + yield ToolCallEvent::Done(result); + return; + } + TaskPayload::Failed { error } => { + yield failed(&task.task_id, TaskErrorReason::Failed { error: serde_json::Value::Object(error) }); + return; + } + TaskPayload::Cancelled => { + yield failed(&task.task_id, TaskErrorReason::Cancelled); + return; + } + _ => { + let reason = TaskErrorReason::UnsupportedPayload { status: task.status }; + yield abandoned(client, &task.task_id, reason).await; return; } } - } - } - - async fn elicit_inputs( - &self, - input_requests: InputRequests, - answered_input_keys: &mut HashSet, - task_id: &str, - ) -> Result<(), TaskErrorReason> { - if input_requests.keys().any(|key| answered_input_keys.contains(key)) { - return Err(TaskErrorReason::RepeatedInput); - } - - let (responses, _) = elicit_inputs(self.client.service(), input_requests).await?; - answered_input_keys.extend(responses.keys().cloned()); - - self.client.update_task(UpdateTaskParams::new(task_id, responses)).await.map_err(TaskErrorReason::Update) - } - - async fn interrupted(&self, task: Task, interrupt: InterruptedReason) -> ToolCallEvent { - match interrupt { - InterruptedReason::TimedOut => self.cancel(task, TaskErrorReason::TimedOut { timeout: self.timeout }).await, - InterruptedReason::Cancelled => { - cancel_server_task(self.client, self.server_name, &task.task_id).await; - ToolCallEvent::Cancelled { task_id: Some(task.task_id) } - } - } - } - - async fn cancel(&self, task: Task, reason: TaskErrorReason) -> ToolCallEvent { - cancel_server_task(self.client, self.server_name, &task.task_id).await; - self.fail(task, reason) - } - - fn fail(&self, task: Task, reason: TaskErrorReason) -> ToolCallEvent { - let error = self.error(&task.task_id, reason); - ToolCallEvent::TaskComplete { task, result: Err(error) } - } - - fn error(&self, task_id: &str, reason: TaskErrorReason) -> CallToolError { - CallToolError::Task { - server: self.server_name.to_string(), - task_id: task_id.to_string(), - reason: Box::new(reason), - } - } -} - -pub(crate) async fn cancel_server_task( - client: &RunningService, - server_name: &str, - task_id: &str, -) { - match timeout(Duration::from_secs(1), client.cancel_task(CancelTaskParams::new(task_id))).await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - tracing::warn!(server = %server_name, %task_id, "Failed to cancel abandoned MCP task: {error}"); - } - Err(_) => tracing::warn!(server = %server_name, %task_id, "Timed out cancelling abandoned MCP task"), - } -} -impl From for TaskErrorReason { - fn from(error: ElicitInputsError) -> Self { - match error { - ElicitInputsError::UnsupportedInput => Self::UnsupportedInput, - ElicitInputsError::Serialize(source) => Self::Serialize(source), + sleep(task.poll_interval_ms.map_or(DEFAULT_POLL_INTERVAL, Duration::from_millis)).await; } } } -struct TaskBounds { - deadline: TaskDeadline, - cancel: CancellationToken, -} +const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(1); -enum InterruptedReason { - TimedOut, - Cancelled, +async fn answer_inputs( + client: &McpClient, + task_id: &str, + input_requests: InputRequests, + answered: &mut HashSet, +) -> Result<(), ToolCallError> { + if input_requests.keys().any(|key| answered.contains(key)) { + return Err(task_error(task_id, TaskErrorReason::RepeatedInput)); + } + let (responses, _) = client.handler.elicit_inputs(input_requests).await?; + answered.extend(responses.keys().cloned()); + client + .peer + .update_task(UpdateTaskParams::new(task_id, responses)) + .await + .map_err(|error| task_error(task_id, TaskErrorReason::Update(error))) } -impl TaskBounds { - fn new(timeout: Duration, cancel: CancellationToken) -> Self { - Self { deadline: TaskDeadline::after(timeout), cancel } - } - - async fn run(&self, future: impl Future) -> Result { - let timedout = pin!(self.deadline.timeout(future)); - let cancelled = pin!(self.cancel.cancelled()); - match select(timedout, cancelled).await { - Either::Left((Ok(value), _)) => Ok(value), - Either::Left((Err(_), _)) => Err(InterruptedReason::TimedOut), - Either::Right(((), _)) => Err(InterruptedReason::Cancelled), - } - } +async fn abandoned(client: &McpClient, task_id: &str, reason: TaskErrorReason) -> ToolCallEvent { + cancel_server_task(client, task_id).await; + failed(task_id, reason) } -enum TaskDeadline { - At(Instant), - FarFuture, +fn failed(task_id: &str, reason: TaskErrorReason) -> ToolCallEvent { + task_error(task_id, reason).into() } -impl TaskDeadline { - fn after(timeout: Duration) -> Self { - Instant::now().checked_add(timeout).map_or(Self::FarFuture, Self::At) - } - - async fn timeout(&self, future: impl Future) -> Result { - match self { - Self::At(deadline) => timeout_at(*deadline, future).await, - Self::FarFuture => Ok(future.await), - } - } +fn task_error(task_id: &str, reason: TaskErrorReason) -> ToolCallError { + ToolCallError::Task { task_id: task_id.to_string(), reason: Box::new(reason) } } fn is_task_expired(task: &Task) -> bool { @@ -303,217 +166,3 @@ fn is_task_expired(task: &Task) -> bool { .checked_add_signed(chrono::Duration::milliseconds(ttl_ms)) .is_some_and(|expires_at| chrono::Utc::now() > expires_at) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::client::call_tool::{CallToolOptions, call_tool}; - use crate::client::{McpClientEvent, client_capabilities}; - use crate::testing::{FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, connect}; - use futures::StreamExt; - use rmcp::model::{ - CallToolRequestParams, CallToolResult, ClientConfig, CreateTaskResult, DetailedTask, ElicitRequest, - ElicitRequestParams, Implementation, InputRequest, ProtocolVersion, - }; - use serde_json::json; - use std::sync::Arc; - use tokio::sync::mpsc; - - #[tokio::test] - async fn call_tool_drives_created_task_to_completion() { - let result = task_test([completed_task()]).run().await; - - assert!( - matches!(result.events.first(), Some(ToolCallEvent::TaskCreated(created)) if created.task.task_id == "task-1") - ); - assert!(matches!( - result.events.last(), - Some(ToolCallEvent::TaskComplete { task, result: Ok(result) }) - if task.task_id == "task-1" - && result.content.first().and_then(|content| content.as_text()).is_some_and(|text| text.text == "finished") - )); - assert_eq!(result.state.task_get_ids(), ["task-1"]); - } - - #[tokio::test] - async fn call_tool_forwards_progress_after_task_creation() { - let seed = task(TaskStatus::Working); - let server = FakeMcpServer::new() - .with_tool( - FakeTool::new("deferred") - .responds(FakeToolResponse::task(CreateTaskResult::new(seed)).task_progress(1.0, Some(2.0))), - ) - .with_task( - "task-1", - [DetailedTask::new(task(TaskStatus::Working), TaskPayload::Working), completed_task()], - ); - let (event_tx, _event_rx) = mpsc::channel::(4); - let client = McpClient::new( - ClientConfig::new(client_capabilities(), Implementation::new("test-client", "0.1.0")), - "task-server".into(), - event_tx, - ); - let (_server, client) = connect(server, client).await.expect("connect task server"); - - let events = call_tool( - Arc::new(client), - CallToolRequestParams::new("deferred"), - CallToolOptions { timeout: Duration::from_secs(1), ..CallToolOptions::default() }, - ) - .collect::>() - .await; - - assert!(matches!(events.first(), Some(ToolCallEvent::TaskCreated(_)))); - assert!(events.iter().any(|event| matches!( - event, - ToolCallEvent::Progress(progress) - if (progress.progress - 1.0).abs() < f64::EPSILON - && progress.total.is_some_and(|total| (total - 2.0).abs() < f64::EPSILON) - ))); - assert!(matches!(events.last(), Some(ToolCallEvent::TaskComplete { result: Ok(_), .. }))); - } - - #[tokio::test] - async fn call_tool_handles_huge_task_ttl() { - let result = - task_test([completed_task()]).with_task(task(TaskStatus::Working).with_ttl_ms(u64::MAX)).run().await; - assert!(matches!(result.events.last(), Some(ToolCallEvent::TaskComplete { result: Ok(_), .. }))); - } - - #[tokio::test] - async fn call_tool_handles_huge_execution_timeout() { - let result = task_test([completed_task()]).with_timeout(Duration::MAX).run().await; - assert!(matches!(result.events.last(), Some(ToolCallEvent::TaskComplete { result: Ok(_), .. }))); - } - - #[tokio::test] - async fn call_tool_cancellation_cancels_server_task_and_ends_stream() { - let seed = task(TaskStatus::Working); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(seed.clone())))) - .with_task("task-1", [DetailedTask::new(seed, TaskPayload::Working)]); - let state = server.state(); - let (event_tx, _event_rx) = mpsc::channel::(4); - let client = McpClient::new( - ClientConfig::new(client_capabilities(), Implementation::new("test-client", "0.1.0")), - "task-server".into(), - event_tx, - ); - let (_server, client) = connect(server, client).await.expect("connect task server"); - - let cancel = CancellationToken::new(); - let options = CallToolOptions { timeout: Duration::from_secs(5), meta: None, cancel: cancel.clone() }; - let mut events = pin!(call_tool(Arc::new(client), CallToolRequestParams::new("deferred"), options)); - - assert!(matches!(events.next().await, Some(ToolCallEvent::TaskCreated(_)))); - cancel.cancel(); - let mut last = None; - while let Some(event) = events.next().await { - last = Some(event); - } - - assert!(matches!(last, Some(ToolCallEvent::Cancelled { task_id: Some(task_id) }) if task_id == "task-1")); - assert_eq!(state.task_cancel_ids(), ["task-1"]); - } - - #[tokio::test] - async fn call_tool_deadline_includes_task_elicitation() { - let result = task_test([input_required_task()]).with_timeout(Duration::from_millis(25)).run().await; - - assert!(matches!( - result.events.last(), - Some(ToolCallEvent::TaskComplete { - result: Err(CallToolError::Task { reason, .. }), - .. - }) if matches!(reason.as_ref(), TaskErrorReason::TimedOut { .. }) - )); - assert_eq!(result.state.task_cancel_ids(), ["task-1"]); - } - - struct TaskTest { - seed: Task, - states: Vec, - timeout: Duration, - } - - struct TaskTestResult { - events: Vec, - state: FakeMcpState, - } - - fn task_test(states: impl IntoIterator) -> TaskTest { - TaskTest { - seed: task(TaskStatus::Working), - states: states.into_iter().collect(), - timeout: Duration::from_secs(1), - } - } - - impl TaskTest { - fn with_task(mut self, seed: Task) -> Self { - self.seed = seed; - self - } - - fn with_timeout(mut self, timeout: Duration) -> Self { - self.timeout = timeout; - self - } - - async fn run(self) -> TaskTestResult { - let task_id = self.seed.task_id.clone(); - let server = FakeMcpServer::new() - .with_tool(FakeTool::new("deferred").responds(FakeToolResponse::task(CreateTaskResult::new(self.seed)))) - .with_task(task_id, self.states); - let state = server.state(); - let (event_tx, _event_rx) = mpsc::channel::(4); - let client = McpClient::new( - ClientConfig::new(client_capabilities(), Implementation::new("test-client", "0.1.0")), - "task-server".into(), - event_tx, - ); - let (_server, client) = connect(server, client).await.expect("connect task server"); - assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2026_07_28); - - let events = call_tool( - Arc::new(client), - CallToolRequestParams::new("deferred"), - CallToolOptions { timeout: self.timeout, ..CallToolOptions::default() }, - ) - .collect() - .await; - TaskTestResult { events, state } - } - } - - fn completed_task() -> DetailedTask { - let result = CallToolResult::success(vec![rmcp::model::ContentBlock::text("finished")]); - DetailedTask::new( - task(TaskStatus::Completed), - TaskPayload::Completed { result: serde_json::from_value(json!(result)).expect("serialize tool result") }, - ) - } - - fn input_required_task() -> DetailedTask { - let request = ElicitRequest::new(ElicitRequestParams::FormElicitationParams { - meta: None, - message: "Provide input".to_string(), - requested_schema: serde_json::from_value(json!({ - "type": "object", - "properties": {} - })) - .expect("valid elicitation schema"), - }); - DetailedTask::new( - task(TaskStatus::InputRequired), - TaskPayload::InputRequired { - input_requests: InputRequests::from([("answer".to_string(), InputRequest::Elicitation(request))]), - }, - ) - } - - fn task(status: TaskStatus) -> Task { - let now = chrono::Utc::now().to_rfc3339(); - Task::new("task-1", status, now.clone(), now).with_poll_interval_ms(10) - } -} diff --git a/crates/mcp-utils/src/client/tool_call.rs b/crates/mcp-utils/src/client/tool_call.rs new file mode 100644 index 000000000..32b61da3c --- /dev/null +++ b/crates/mcp-utils/src/client/tool_call.rs @@ -0,0 +1,329 @@ +use super::McpClient; +use super::task::{TaskErrorReason, cancel_server_task, stream_task_events}; +use async_stream::stream; +use futures::future::Either; +use futures::{Stream, StreamExt, stream}; +use rmcp::RoleClient; +use rmcp::model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, ClientRequest, CreateTaskResult, DEFAULT_MRTR_MAX_ROUNDS, + InputRequests, InputRequiredResult, ProgressNotificationParam, Request, RequestMetaObject, ServerResult, Task, +}; +use rmcp::service::{PeerRequestOptions, RequestHandle, ServiceError}; +use std::future::pending; +use std::pin::{Pin, pin}; +use std::task::{Context, Poll}; +use std::time::Duration; +use thiserror::Error; +use tokio::time::{Instant, sleep, sleep_until}; +use tokio_util::sync::CancellationToken; + +pub struct ToolCall { + events: Pin + Send>>, +} + +#[derive(Debug)] +pub enum ToolCallEvent { + Progress(ProgressNotificationParam), + TaskCreated(CreateTaskResult), + TaskStatus(Task), + Done(Result), +} + +#[derive(Debug, Clone, Default)] +pub struct ToolCallOptions { + pub timeout: Option, + pub meta: Option, + pub cancel: CancellationToken, +} + +#[derive(Debug, Error)] +pub enum ToolCallError { + #[error("Failed to send tool request: {0}")] + Send(#[source] ServiceError), + #[error("Tool execution failed: {0}")] + Call(#[source] ServiceError), + #[error("Tool call exceeded its {0:?} timeout")] + TimedOut(Duration), + #[error("Server requested an input kind this client does not support (sampling or roots)")] + UnsupportedInput, + #[error("Server requested input without any input requests or request state")] + EmptyInputRequired, + #[error("Server requested input again after the user cancelled")] + RePromptAfterCancel, + #[error("Server did not complete within {DEFAULT_MRTR_MAX_ROUNDS} MRTR input rounds")] + InputRoundsExceeded, + #[error("Task '{task_id}' {reason}")] + Task { + task_id: String, + #[source] + reason: Box, + }, + #[error("Server returned a tool call response kind this client does not support")] + UnsupportedResponse, + #[error("Tool call was cancelled")] + Cancelled, +} + +impl ToolCallOptions { + pub fn with_timeout(timeout: Duration) -> Self { + Self { timeout: Some(timeout), ..Self::default() } + } +} + +impl ToolCall { + pub(crate) fn new(client: McpClient, params: CallToolRequestParams, options: ToolCallOptions) -> Self { + let ToolCallOptions { timeout, meta, cancel } = options; + let events = stream_events(client.clone(), params, meta); + Self { events: Box::pin(bounded(client, events, timeout, cancel)) } + } + + pub async fn result(self) -> Result { + self.result_with_progress(|_| std::future::ready(())).await + } + + pub async fn result_with_progress>( + mut self, + mut on_progress: impl FnMut(ProgressNotificationParam) -> F, + ) -> Result { + while let Some(event) = self.next().await { + match event { + ToolCallEvent::Progress(progress) => on_progress(progress).await, + ToolCallEvent::Done(result) => return result, + ToolCallEvent::TaskCreated(_) | ToolCallEvent::TaskStatus(_) => {} + } + } + Err(ToolCallError::Cancelled) + } +} + +impl Stream for ToolCall { + type Item = ToolCallEvent; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.events.as_mut().poll_next(cx) + } +} + +impl From for ToolCallEvent { + fn from(error: ToolCallError) -> Self { + Self::Done(Err(error)) + } +} + +pub(super) fn with_progress<'a, T: 'a>( + progress: impl Stream + 'a, + events: impl Stream + 'a, +) -> impl Stream> + 'a { + stream! { + let mut progress = pin!(progress); + let mut events = pin!(events); + loop { + tokio::select! { + biased; + progress_event = progress.next() => match progress_event { + Some(progress_event) => yield Either::Left(progress_event), + None => { + while let Some(event) = events.next().await { + yield Either::Right(event); + } + return; + } + }, + event = events.next() => match event { + Some(event) => yield Either::Right(event), + None => return, + }, + } + } + } +} + +fn bounded( + client: McpClient, + events: impl Stream + Send + 'static, + timeout: Option, + cancel: CancellationToken, +) -> impl Stream + Send + 'static { + stream! { + let mut events = Box::pin(events); + let mut expired = pin!(expiry(timeout)); + let mut cancelled = pin!(cancel.cancelled()); + let mut task = None; + + let interrupted = loop { + tokio::select! { + biased; + event = events.next() => { + let Some(event) = event else { return }; + match &event { + ToolCallEvent::TaskCreated(created) => task = Some(created.task.clone()), + ToolCallEvent::TaskStatus(status) => task = Some(status.clone()), + ToolCallEvent::Progress(_) | ToolCallEvent::Done(_) => {} + } + let done = matches!(event, ToolCallEvent::Done(_)); + yield event; + if done { + return; + } + } + timeout = &mut expired => break ToolCallError::TimedOut(timeout), + () = &mut cancelled => break ToolCallError::Cancelled, + } + }; + + drop(events); + if let Some(task) = &task { + cancel_server_task(&client, &task.task_id).await; + } + yield interrupted.into(); + } +} + +fn stream_events( + client: McpClient, + mut params: CallToolRequestParams, + meta: Option, +) -> impl Stream + Send + 'static { + stream! { + let mut mrtr_state = MrtrState::new(); + + loop { + let request = ClientRequest::CallToolRequest(Request::new(params.clone())); + let handle = match client.peer.send_cancellable_request(request, peer_request_options(meta.clone())).await { + Ok(handle) => handle, + Err(error) => { + yield ToolCallError::Send(error).into(); + return; + } + }; + + let mut progress = client.handler.progress.subscribe(handle.progress_token.clone()).await; + let mut response = None; + { + let mut awaited = pin!(with_progress(&mut progress, stream::once(await_tool_response(handle)))); + while let Some(event) = awaited.next().await { + match event { + Either::Left(progress_event) => yield ToolCallEvent::Progress(progress_event), + Either::Right(awaited_response) => response = Some(awaited_response), + } + } + } + + match response.unwrap_or(Err(ServiceError::UnexpectedResponse)) { + Ok(CallToolResponse::Complete(result)) => { + yield ToolCallEvent::Done(Ok(result)); + return; + } + Ok(CallToolResponse::InputRequired(input_required)) => match mrtr_state.tick(input_required) { + MrtrAction::Poll { backoff, request_state } => { + sleep(backoff).await; + params.input_responses = None; + params.request_state = Some(request_state); + } + MrtrAction::Elicit { input_requests, request_state } => { + let (responses, cancelled) = match client.handler.elicit_inputs(input_requests).await { + Ok(answered) => answered, + Err(error) => { + yield error.into(); + return; + } + }; + mrtr_state.record_cancelled(cancelled); + params.input_responses = Some(responses); + params.request_state = request_state; + } + MrtrAction::Abort(error) => { + yield error.into(); + return; + } + }, + Ok(CallToolResponse::Task(created)) => { + let mut events = pin!(stream_task_events(&client, created, progress)); + while let Some(event) = events.next().await { + yield event; + } + return; + } + Ok(_) => { + yield ToolCallError::UnsupportedResponse.into(); + return; + } + Err(error) => { + yield ToolCallError::Call(error).into(); + return; + } + } + } + } +} + +async fn expiry(timeout: Option) -> Duration { + match timeout.and_then(|timeout| Some((Instant::now().checked_add(timeout)?, timeout))) { + Some((deadline, timeout)) => { + sleep_until(deadline).await; + timeout + } + None => pending().await, + } +} + +fn peer_request_options(meta: Option) -> PeerRequestOptions { + meta.map_or_else(PeerRequestOptions::no_options, |meta| PeerRequestOptions::no_options().with_meta(meta)) +} + +async fn await_tool_response(handle: RequestHandle) -> Result { + match handle.await_response().await? { + ServerResult::CallToolResult(result) => Ok(CallToolResponse::Complete(result)), + ServerResult::InputRequiredResult(result) => Ok(CallToolResponse::InputRequired(result)), + ServerResult::CreateTaskResult(result) => Ok(CallToolResponse::Task(result)), + _ => Err(ServiceError::UnexpectedResponse), + } +} + +struct MrtrState { + input_request_rounds: usize, + next_backoff: Duration, + user_cancelled: bool, +} + +enum MrtrAction { + Poll { backoff: Duration, request_state: String }, + Elicit { input_requests: InputRequests, request_state: Option }, + Abort(ToolCallError), +} + +impl MrtrState { + fn new() -> Self { + Self { input_request_rounds: 0, next_backoff: BASE_BACKOFF, user_cancelled: false } + } + + fn tick(&mut self, input_required: InputRequiredResult) -> MrtrAction { + let input_requests = input_required.input_requests.filter(|requests| !requests.is_empty()); + match (input_requests, input_required.request_state) { + (None, None) => MrtrAction::Abort(ToolCallError::EmptyInputRequired), + (None, Some(request_state)) => { + let backoff = self.next_backoff; + self.next_backoff = (backoff * 2).min(MAX_BACKOFF); + MrtrAction::Poll { backoff, request_state } + } + (Some(input_requests), request_state) => { + if self.user_cancelled { + MrtrAction::Abort(ToolCallError::RePromptAfterCancel) + } else if self.input_request_rounds == DEFAULT_MRTR_MAX_ROUNDS { + MrtrAction::Abort(ToolCallError::InputRoundsExceeded) + } else { + self.input_request_rounds += 1; + self.next_backoff = BASE_BACKOFF; + MrtrAction::Elicit { input_requests, request_state } + } + } + } + } + + fn record_cancelled(&mut self, cancelled: bool) { + self.user_cancelled |= cancelled; + } +} + +const BASE_BACKOFF: Duration = Duration::from_millis(50); +const MAX_BACKOFF: Duration = Duration::from_millis(1600); diff --git a/crates/mcp-utils/src/client/tool_catalog.rs b/crates/mcp-utils/src/client/tool_catalog.rs deleted file mode 100644 index 4ab3fadf8..000000000 --- a/crates/mcp-utils/src/client/tool_catalog.rs +++ /dev/null @@ -1,295 +0,0 @@ -use super::{ToolExposure, connection::Tool, naming::create_namespaced_tool_name, tool_filter::ToolFilter}; -use llm::ToolDefinition; -use std::collections::BTreeMap; -use utils::mcp_status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry}; - -pub const PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME: &str = "progressive-discovery"; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ServerDescription { - pub name: String, - pub description: String, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct ToolCatalog { - servers: Vec, - progressive_discovery_instructions: Option, -} - -#[derive(Debug, Clone, PartialEq)] -pub struct ServerCatalogEntry { - name: String, - description: String, - instructions: Option, - status: McpServerStatus, - auth_capability: McpServerAuthCapability, - exposure: ToolExposure, - tools: Vec, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct CatalogTools<'a> { - pub model_visible: Vec<&'a CatalogTool>, - pub deferred: Vec<&'a CatalogTool>, -} - -#[derive(Debug, Clone, PartialEq)] -pub struct CatalogTool { - namespaced_name: String, - local_name: String, - definition: ToolDefinition, - exposure: ToolExposureKind, - allowed: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ToolExposureKind { - ModelVisible, - Deferred, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum ToolRoute { - ModelVisible { namespaced_name: String }, - Deferred { server: String, tool: String }, -} - -impl ToolCatalog { - pub fn new() -> Self { - Self::default() - } - pub fn servers(&self) -> &[ServerCatalogEntry] { - &self.servers - } - - pub fn server(&self, name: &str) -> Option<&ServerCatalogEntry> { - self.servers.iter().find(|server| server.name == name) - } - - pub fn tool(&self, namespaced_name: &str) -> Option<&CatalogTool> { - self.servers.iter().flat_map(|server| &server.tools).find(|tool| tool.namespaced_name == namespaced_name) - } - - pub fn tools(&self) -> CatalogTools<'_> { - CatalogTools::from_tools( - self.servers.iter().filter(|server| server.is_connected()).flat_map(|server| server.tools.iter()), - ) - } - - pub fn tools_for(&self, server: &str) -> Option> { - self.server(server).map(|entry| { - if entry.is_connected() { CatalogTools::from_tools(entry.tools.iter()) } else { CatalogTools::default() } - }) - } - - pub fn discoverable_deferred_servers(&self) -> Vec { - self.servers - .iter() - .filter(|server| { - server.is_connected() - && server.tools.iter().any(|tool| tool.allowed && tool.exposure == ToolExposureKind::Deferred) - }) - .map(|server| ServerDescription { name: server.name.clone(), description: server.description.clone() }) - .collect() - } - - pub fn model_instructions(&self) -> BTreeMap { - let mut instructions = self - .servers - .iter() - .filter(|server| server.is_connected()) - .filter(|server| { - server.tools.iter().any(|tool| tool.allowed && tool.exposure == ToolExposureKind::ModelVisible) - }) - .filter_map(|server| server.instructions.as_ref().map(|body| (server.name.clone(), body.clone()))) - .collect::>(); - if !self.discoverable_deferred_servers().is_empty() - && let Some(body) = &self.progressive_discovery_instructions - { - instructions.insert(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME.to_string(), body.clone()); - } - instructions - } - - pub fn route_permitted(&self, route: &ToolRoute) -> bool { - let (namespaced_name, exposure) = match route { - ToolRoute::ModelVisible { namespaced_name } => (namespaced_name.clone(), ToolExposureKind::ModelVisible), - ToolRoute::Deferred { server, tool } => { - (create_namespaced_tool_name(server, tool), ToolExposureKind::Deferred) - } - }; - let Some(server) = self.servers.iter().find(|server| { - server.is_connected() && server.tools.iter().any(|tool| tool.namespaced_name == namespaced_name) - }) else { - return false; - }; - let Some(tool) = server.tools.iter().find(|tool| tool.namespaced_name == namespaced_name) else { return false }; - tool.allowed && tool.exposure == exposure - } - - pub fn server_statuses(&self) -> Vec { - self.servers.iter().map(ServerCatalogEntry::status_entry).collect() - } - - pub fn upsert_server(&mut self, entry: ServerCatalogEntry) { - if let Some(existing) = self.servers.iter_mut().find(|server| server.name == entry.name) { - *existing = entry; - } else { - self.servers.push(entry); - } - } - - pub fn remove_server(&mut self, name: &str) -> Option { - self.servers.iter().position(|server| server.name == name).map(|index| self.servers.remove(index)) - } - - pub fn set_progressive_discovery_instructions(&mut self, instructions: Option) { - self.progressive_discovery_instructions = instructions; - } -} - -impl ServerCatalogEntry { - #[allow(clippy::too_many_arguments)] - pub fn new( - name: impl Into, - description: impl Into, - instructions: Option, - status: McpServerStatus, - auth_capability: McpServerAuthCapability, - exposure: ToolExposure, - tools: &[rmcp::model::Tool], - filter: &ToolFilter, - ) -> Self { - let tools = tools.iter().map(Tool::from).collect::>(); - Self::from_tools( - name.into(), - description.into(), - instructions, - status, - auth_capability, - exposure, - &tools, - filter, - ) - } - pub fn name(&self) -> &str { - &self.name - } - pub fn description(&self) -> &str { - &self.description - } - pub fn instructions(&self) -> Option<&str> { - self.instructions.as_deref() - } - pub fn status(&self) -> &McpServerStatus { - &self.status - } - pub fn auth_capability(&self) -> McpServerAuthCapability { - self.auth_capability - } - pub fn exposure(&self) -> &ToolExposure { - &self.exposure - } - pub fn tools(&self) -> &[CatalogTool] { - &self.tools - } - pub fn status_entry(&self) -> McpServerStatusEntry { - McpServerStatusEntry::new(&self.name, self.status.clone()) - .with_auth_capability(self.auth_capability) - .with_deferred_tools(self.exposure.has_deferred_tools()) - } - pub(crate) fn pending(name: impl Into, exposure: ToolExposure) -> Self { - let name = name.into(); - Self { - description: name.clone(), - name, - instructions: None, - status: McpServerStatus::Connecting, - auth_capability: McpServerAuthCapability::Unavailable, - exposure, - tools: Vec::new(), - } - } - #[allow(clippy::too_many_arguments)] - pub(crate) fn from_tools( - name: String, - description: String, - instructions: Option, - status: McpServerStatus, - auth_capability: McpServerAuthCapability, - exposure: ToolExposure, - tools: &[Tool], - filter: &ToolFilter, - ) -> Self { - let catalog_tools = tools - .iter() - .map(|tool| { - let definition = ToolDefinition::new( - create_namespaced_tool_name(&name, &tool.name), - tool.description.clone(), - tool.parameters.clone(), - ) - .with_server(name.clone()) - .with_annotations(tool.annotations.clone()); - let exposure_kind = if exposure.is_model_visible_tool(&tool.name) { - ToolExposureKind::ModelVisible - } else { - ToolExposureKind::Deferred - }; - CatalogTool { - namespaced_name: definition.name.clone(), - local_name: tool.name.clone(), - allowed: filter.is_tool_allowed(&definition), - definition, - exposure: exposure_kind, - } - }) - .collect(); - Self { name, description, instructions, status, auth_capability, exposure, tools: catalog_tools } - } - pub(crate) fn with_status(&self, status: McpServerStatus, auth_capability: McpServerAuthCapability) -> Self { - let mut next = self.clone(); - next.status = status; - next.auth_capability = auth_capability; - if !next.is_connected() { - next.tools.clear(); - next.instructions = None; - } - next - } - fn is_connected(&self) -> bool { - matches!(self.status, McpServerStatus::Connected { .. }) - } -} - -impl<'a> CatalogTools<'a> { - fn from_tools(tools: impl Iterator) -> Self { - let mut partitioned = Self::default(); - for tool in tools.filter(|tool| tool.allowed) { - match tool.exposure { - ToolExposureKind::ModelVisible => partitioned.model_visible.push(tool), - ToolExposureKind::Deferred => partitioned.deferred.push(tool), - } - } - partitioned - } -} - -impl CatalogTool { - pub fn namespaced_name(&self) -> &str { - &self.namespaced_name - } - pub fn local_name(&self) -> &str { - &self.local_name - } - pub fn definition(&self) -> &ToolDefinition { - &self.definition - } - pub fn exposure(&self) -> ToolExposureKind { - self.exposure - } - pub fn allowed(&self) -> bool { - self.allowed - } -} diff --git a/crates/mcp-utils/src/client/transport.rs b/crates/mcp-utils/src/client/transport.rs new file mode 100644 index 000000000..ee823bb07 --- /dev/null +++ b/crates/mcp-utils/src/client/transport.rs @@ -0,0 +1,176 @@ +use super::error::{ClientError, Result}; +use super::{ + ClientOptions, + handler::Handler, + oauth::{self, OAuthSettings}, +}; +use crate::server::McpServer; +use reqwest::header::HeaderMap; +use rmcp::{ + ClientLifecycleMode, RoleClient, + model::ProtocolVersion, + serve_client_with_lifecycle, + service::{ClientInitializeError, RunningService}, + transport::{ + IntoTransport, StreamableHttpClientTransport, TokioChildProcess, + streamable_http_client::StreamableHttpClientTransportConfig, + }, +}; +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Arc; +use tokio::{ + io::{AsyncBufReadExt, BufReader}, + net::UnixStream, + process::{ChildStderr, Command}, +}; + +#[derive(Debug, Clone)] +pub enum Transport { + Stdio { command: String, args: Vec, env: HashMap, cwd: Option }, + Http { url: String, headers: HeaderMap, oauth: Option }, + InProcess(McpServer), + Unix(PathBuf), +} + +pub(super) type ClientService = RunningService>; + +impl Transport { + pub(crate) fn supports_oauth(&self, options: &ClientOptions) -> bool { + options.oauth_prompts().is_some() && matches!(self, Self::Http { oauth: Some(_), .. }) + } + + pub(super) async fn connect(self, handler: Handler, options: &ClientOptions) -> Result { + let supports_oauth = self.supports_oauth(options); + let handler = Arc::new(handler); + let connected = match self { + Self::Stdio { command, args, env, cwd } => connect_stdio(handler, command, args, env, cwd).await, + Self::Http { url, headers, oauth } => connect_http(handler, url, headers, oauth, options).await, + Self::InProcess(server) => serve(handler, server.serve_in_memory()).await, + Self::Unix(path) => connect_unix(handler, path).await, + }; + match connected { + Err(ClientError::Connect { server, source }) if supports_oauth && requires_authorization(&source) => { + Err(ClientError::AuthRequired { server, challenge: source.auth_challenge().map(str::to_string) }) + } + connected => connected, + } + } + + pub(super) async fn authorize( + self, + handler: Handler, + options: &ClientOptions, + challenge: Option, + ) -> Result { + let server = handler.server().to_string(); + let Self::Http { url, headers, oauth: Some(oauth) } = self else { + return Err(ClientError::OAuthUnavailable { server }); + }; + let auth_client = oauth::authorize(&server, &url, oauth, options, challenge).await?; + let transport = StreamableHttpClientTransport::with_client(auth_client, http_config(url, headers)); + serve(Arc::new(handler), transport).await + } +} + +fn client_lifecycle_mode() -> ClientLifecycleMode { + ClientLifecycleMode::Auto { + preferred_versions: vec![ + ProtocolVersion::V_2026_07_28, + ProtocolVersion::V_2025_11_25, + ProtocolVersion::V_2025_06_18, + ProtocolVersion::V_2025_03_26, + ProtocolVersion::V_2024_11_05, + ], + legacy_version: Some(ProtocolVersion::V_2025_11_25), + } +} + +async fn connect_stdio( + handler: Arc, + command: String, + args: Vec, + env: HashMap, + cwd: Option, +) -> Result { + let mut cmd = Command::new(&command); + cmd.args(&args).envs(&env); + if let Some(cwd) = cwd { + cmd.current_dir(cwd); + } + let (process, stderr) = TokioChildProcess::builder(cmd) + .stderr(Stdio::piped()) + .spawn() + .map_err(|source| ClientError::Spawn { command, source })?; + + if let Some(stderr) = stderr { + spawn_stderr_logger(handler.server().to_string(), stderr); + } + + serve(handler, process).await +} + +async fn connect_http( + handler: Arc, + url: String, + headers: HeaderMap, + oauth: Option, + options: &ClientOptions, +) -> Result { + let server = handler.server(); + let restored = match (options.oauth_store.as_ref(), &oauth) { + (Some(store), Some(oauth)) => oauth::restore(server, &url, oauth, store).await, + _ => None, + }; + + let config = http_config(url, headers); + match restored { + Some(auth_client) => { + tracing::debug!(server, "Using stored OAuth credentials"); + serve(handler, StreamableHttpClientTransport::with_client(auth_client, config)).await + } + None => serve(handler, StreamableHttpClientTransport::from_config(config)).await, + } +} + +async fn connect_unix(handler: Arc, path: PathBuf) -> Result { + let stream = UnixStream::connect(&path).await.map_err(|source| ClientError::SocketConnect { path, source })?; + serve(handler, stream).await +} + +async fn serve(handler: Arc, transport: T) -> Result +where + T: IntoTransport, + U: std::error::Error + Send + Sync + 'static, +{ + let server = handler.server().to_string(); + serve_client_with_lifecycle(handler, transport, client_lifecycle_mode()) + .await + .map_err(|source| ClientError::connect(server, source)) +} + +fn requires_authorization(error: &ClientInitializeError) -> bool { + error.is_authorization_required() || error.auth_challenge().is_some() +} + +fn http_config(url: String, headers: HeaderMap) -> StreamableHttpClientTransportConfig { + let headers = headers.into_iter().filter_map(|(name, value)| Some((name?, value))).collect(); + StreamableHttpClientTransportConfig::with_uri(url).custom_headers(headers) +} + +fn spawn_stderr_logger(server: String, stderr: ChildStderr) { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => tracing::info!(server = %server, stderr = %line, "MCP server stderr"), + Ok(None) => break, + Err(error) => { + tracing::warn!(server = %server, %error, "failed to read MCP server stderr"); + break; + } + } + } + }); +} diff --git a/crates/mcp-utils/src/gateway/connection.rs b/crates/mcp-utils/src/gateway/connection.rs new file mode 100644 index 000000000..b916d0dac --- /dev/null +++ b/crates/mcp-utils/src/gateway/connection.rs @@ -0,0 +1,225 @@ +use super::{ConnectedServer, GatewayError, ServerSpec, ServerState, ServerStatus}; +use crate::client::{ClientError, ClientOptions, McpClient}; +use aether_auth::OAuthError; +use async_stream::stream; +use futures::{Stream, StreamExt}; +use std::pin::pin; +use tokio::sync::{mpsc, oneshot, watch}; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +pub(super) struct Connection { + name: String, + supports_oauth: bool, + state: watch::Receiver, + sign_ins: mpsc::UnboundedSender, +} + +#[derive(Debug, Clone)] +pub(super) enum ConnectionState { + Connecting, + Authenticating, + NeedsOAuth { challenge: Option }, + Failed { error: String, after_sign_in: bool }, + Connected(ConnectedServer), +} + +impl Connection { + pub(super) fn spawn( + spec: ServerSpec, + options: &ClientOptions, + stop: CancellationToken, + tasks: &TaskTracker, + ) -> Self { + let (state, receiver) = watch::channel(ConnectionState::Connecting); + let (sign_ins, requests) = mpsc::unbounded_channel(); + let supports_oauth = spec.transport.supports_oauth(options); + let connection = Self { name: spec.name.clone(), supports_oauth, state: receiver, sign_ins }; + let task = ConnectionTask { spec, options: options.clone(), supports_oauth, state }; + tasks.spawn(task.run(requests, stop)); + connection + } + + pub(super) fn name(&self) -> &str { + &self.name + } + + pub(super) fn connected(&self) -> Option { + self.state.borrow().connected().cloned() + } + + pub(super) fn subscribe(&self) -> watch::Receiver { + let mut receiver = self.state.clone(); + receiver.mark_unchanged(); + receiver + } + + pub(super) async fn settled(&self) { + let _ = self.subscribe().wait_for(|state| !matches!(state, ConnectionState::Connecting)).await; + } + + pub(super) async fn sign_in(&self) -> Result<(), GatewayError> { + let (reply, outcome) = oneshot::channel(); + self.sign_ins.send(reply).map_err(|_| GatewayError::Closed)?; + outcome.await.unwrap_or(Err(GatewayError::Closed)) + } + + pub(super) fn status(&self) -> ServerStatus { + let state = self.state.borrow(); + ServerStatus { + name: self.name.clone(), + state: state.status(), + can_sign_in: state.allows_sign_in(self.supports_oauth), + } + } +} + +impl ConnectionState { + fn connected(&self) -> Option<&ConnectedServer> { + match self { + Self::Connected(server) => Some(server), + _ => None, + } + } + + fn allows_sign_in(&self, supports_oauth: bool) -> bool { + supports_oauth + && match self { + Self::Connecting => false, + Self::Failed { after_sign_in, .. } => *after_sign_in, + Self::Authenticating | Self::NeedsOAuth { .. } | Self::Connected(_) => true, + } + } + + fn status(&self) -> ServerState { + match self { + Self::Connecting => ServerState::Connecting, + Self::Authenticating => ServerState::Authenticating, + Self::NeedsOAuth { .. } => ServerState::NeedsOAuth, + Self::Failed { error, .. } => ServerState::Failed { error: error.clone() }, + Self::Connected(server) => ServerState::Connected { tool_count: server.tools().len() }, + } + } +} + +type Reply = oneshot::Sender>; + +type Update = Result; + +enum Credentials { + Existing, + SignIn { challenge: Option }, +} + +struct ConnectionTask { + spec: ServerSpec, + options: ClientOptions, + supports_oauth: bool, + state: watch::Sender, +} + +impl ConnectionTask { + async fn run(self, mut sign_ins: mpsc::UnboundedReceiver, stop: CancellationToken) { + let mut pending = None; + let mut connection = pin!(self.connection(Credentials::Existing)); + loop { + tokio::select! { + () = stop.cancelled() => break, + Some(update) = connection.next() => self.settle(update, &mut pending), + Some(reply) = sign_ins.recv() => { + if let Some(credentials) = self.sign_in(reply, &mut pending) { + connection.set(self.connection(credentials)); + } + } + } + } + let server = self.state.borrow().connected().cloned(); + if let Some(server) = server { + server.client().close().await; + } + } + + fn connection(&self, credentials: Credentials) -> impl Stream + Send + '_ { + stream! { + let (name, transport) = (&self.spec.name, self.spec.transport.clone()); + let connected = match credentials { + Credentials::Existing => McpClient::connect(name, transport, &self.options).await, + Credentials::SignIn { challenge } => McpClient::authorize(name, transport, &self.options, challenge).await, + }; + let client = match connected { + Ok(client) => client, + Err(error) => { + yield Err(error.into()); + return; + } + }; + let _close = client.close_on_drop(); + let mut tool_lists = pin!(client.tool_lists()); + let mut current: Option = None; + while let Some(tools) = tool_lists.next().await { + match (tools, ¤t) { + (Ok(tools), Some(server)) if server.tools() == tools.as_slice() => {} + (Ok(tools), _) => { + let server = ConnectedServer::new(client.clone(), tools); + current = Some(server.clone()); + yield Ok(server); + } + (Err(error), None) => { + yield Err(error.into()); + return; + } + (Err(error), Some(_)) => { + tracing::warn!(server = name, %error, "Failed to refresh MCP tools; keeping the previous list"); + } + } + } + yield Err(ClientError::Disconnected { server: name.clone() }.into()); + } + } + + fn sign_in(&self, reply: Reply, pending: &mut Option) -> Option { + let server = self.spec.name.clone(); + if !self.can_sign_in() { + let _ = reply.send(Err(ClientError::OAuthUnavailable { server }.into())); + return None; + } + let challenge = match self.state.send_replace(ConnectionState::Authenticating) { + ConnectionState::NeedsOAuth { challenge } => challenge, + _ => None, + }; + if let Some(replaced) = pending.replace(reply) { + let _ = replaced.send(Err(ClientError::Auth { server, source: OAuthError::UserCancelled }.into())); + } + Some(Credentials::SignIn { challenge }) + } + + fn settle(&self, update: Update, pending: &mut Option) { + let (state, result) = match update { + Ok(server) => (ConnectionState::Connected(server), Ok(())), + Err(error) => (self.failed(&error), Err(error)), + }; + self.state.send_replace(state); + if let Some(reply) = pending.take() { + let _ = reply.send(result); + } + } + + fn failed(&self, error: &GatewayError) -> ConnectionState { + let server = &self.spec.name; + match error { + GatewayError::Client(ClientError::AuthRequired { challenge, .. }) => { + tracing::warn!(server, "MCP server needs OAuth"); + ConnectionState::NeedsOAuth { challenge: challenge.clone() } + } + error => { + tracing::warn!(server, %error, "MCP server failed"); + let after_sign_in = matches!(*self.state.borrow(), ConnectionState::Authenticating); + ConnectionState::Failed { error: error.to_string(), after_sign_in } + } + } + } + + fn can_sign_in(&self) -> bool { + self.state.borrow().allows_sign_in(self.supports_oauth) + } +} diff --git a/crates/mcp-utils/src/gateway/error.rs b/crates/mcp-utils/src/gateway/error.rs new file mode 100644 index 000000000..88bfa00cc --- /dev/null +++ b/crates/mcp-utils/src/gateway/error.rs @@ -0,0 +1,12 @@ +use crate::client::ClientError; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum GatewayError { + #[error("Server not found: {0}")] + ServerNotFound(String), + #[error("MCP gateway is closed")] + Closed, + #[error(transparent)] + Client(#[from] ClientError), +} diff --git a/crates/mcp-utils/src/gateway/mcp_gateway.rs b/crates/mcp-utils/src/gateway/mcp_gateway.rs new file mode 100644 index 000000000..6da0bef7c --- /dev/null +++ b/crates/mcp-utils/src/gateway/mcp_gateway.rs @@ -0,0 +1,172 @@ +use super::GatewayError; +use super::connection::Connection; +use crate::client::{ClientOptions, McpClient, ToolCall, ToolCallOptions, Transport}; +use crate::server::McpServer; +use futures::future::join_all; +use futures::stream::select_all; +use futures::{Stream, StreamExt}; +use rmcp::model::{GetPromptResult, Prompt, Tool}; +use serde_json::{Map, Value}; +use std::sync::Arc; +use tokio_stream::wrappers::WatchStream; +use tokio_util::sync::{CancellationToken, DropGuard}; +use tokio_util::task::TaskTracker; + +#[derive(Clone)] +pub struct McpGateway { + connections: Arc<[Connection]>, + stop: CancellationToken, + tasks: TaskTracker, + _stop_on_drop: Arc, +} + +#[derive(Debug, Clone)] +pub struct ServerSpec { + pub name: String, + pub transport: Transport, +} + +#[derive(Debug, Clone)] +pub struct ConnectedServer { + client: McpClient, + tools: Arc<[Tool]>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ServerStatus { + pub name: String, + pub state: ServerState, + pub can_sign_in: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ServerState { + Connecting, + Authenticating, + NeedsOAuth, + Failed { error: String }, + Connected { tool_count: usize }, +} + +#[derive(Debug, Clone)] +pub struct ServerPrompt { + pub server: String, + pub prompt: Prompt, +} + +impl McpGateway { + /// Connects to every server in `specs`, which must have unique names. + pub fn new(specs: Vec, options: &ClientOptions) -> Self { + let stop = CancellationToken::new(); + let tasks = TaskTracker::new(); + let connections = + specs.into_iter().map(|spec| Connection::spawn(spec, options, stop.clone(), &tasks)).collect(); + + tasks.close(); + Self { connections, _stop_on_drop: Arc::new(stop.clone().drop_guard()), stop, tasks } + } + + pub async fn ready(&self) { + join_all(self.connections().iter().map(Connection::settled)).await; + } + + pub fn changes(&self) -> impl Stream + Send + Sync + Unpin + use<> { + select_all(self.connections.iter().map(|connection| WatchStream::from_changes(connection.subscribe()))) + .map(|_| ()) + } + + pub fn servers(&self) -> Vec { + self.connections().iter().filter_map(Connection::connected).collect() + } + + pub fn statuses(&self) -> Vec { + self.connections().iter().map(Connection::status).collect() + } + + pub async fn list_prompts(&self) -> Vec { + let listed = join_all(self.servers().into_iter().map(|server| async move { + match server.client.list_prompts().await { + Ok(prompts) => prompts + .into_iter() + .map(|prompt| ServerPrompt { server: server.name().to_string(), prompt }) + .collect(), + Err(error) => { + tracing::warn!(server = server.name(), %error, "Failed to list MCP prompts"); + Vec::new() + } + } + })) + .await; + listed.into_iter().flatten().collect() + } + + pub async fn get_prompt( + &self, + server: &str, + prompt: &str, + args: Map, + ) -> Result { + let connected = self.server(server).ok_or_else(|| GatewayError::ServerNotFound(server.to_string()))?; + Ok(connected.client.get_prompt(prompt, args).await?) + } + + pub async fn authenticate(&self, server: &str) -> Result<(), GatewayError> { + if self.stop.is_cancelled() { + return Err(GatewayError::Closed); + } + self.connection(server).ok_or_else(|| GatewayError::ServerNotFound(server.to_string()))?.sign_in().await + } + + pub async fn shutdown(&self) { + self.stop.cancel(); + self.tasks.wait().await; + } + + pub fn server(&self, name: &str) -> Option { + self.connection(name)?.connected() + } + + fn connections(&self) -> &[Connection] { + if self.stop.is_cancelled() { &[] } else { &self.connections } + } + + fn connection(&self, name: &str) -> Option<&Connection> { + self.connections().iter().find(|connection| connection.name() == name) + } +} + +impl ServerSpec { + pub fn in_process(name: impl Into, server: impl Into) -> Self { + Self { name: name.into(), transport: Transport::InProcess(server.into()) } + } +} + +impl ConnectedServer { + pub fn name(&self) -> &str { + self.client.name() + } + + pub fn description(&self) -> Option { + self.client.description() + } + + pub fn instructions(&self) -> Option { + self.client.instructions() + } + + pub fn tools(&self) -> &[Tool] { + &self.tools + } + + pub fn call_tool(&self, name: &str, args: Map, options: ToolCallOptions) -> ToolCall { + self.client.call_tool(name, args, options) + } + + pub(super) fn new(client: McpClient, tools: Vec) -> Self { + Self { client, tools: tools.into() } + } + + pub(super) fn client(&self) -> &McpClient { + &self.client + } +} diff --git a/crates/mcp-utils/src/gateway/mod.rs b/crates/mcp-utils/src/gateway/mod.rs new file mode 100644 index 000000000..b5e614628 --- /dev/null +++ b/crates/mcp-utils/src/gateway/mod.rs @@ -0,0 +1,6 @@ +mod connection; +mod error; +mod mcp_gateway; + +pub use error::GatewayError; +pub use mcp_gateway::{ConnectedServer, McpGateway, ServerPrompt, ServerSpec, ServerState, ServerStatus}; diff --git a/crates/mcp-utils/src/lib.rs b/crates/mcp-utils/src/lib.rs index d9a25e0e3..863b31140 100644 --- a/crates/mcp-utils/src/lib.rs +++ b/crates/mcp-utils/src/lib.rs @@ -1,11 +1,11 @@ #![doc = include_str!("../README.md")] -mod protocol; -pub mod server; -pub mod testing; -pub mod tool_gateway; - #[cfg(feature = "client")] pub mod client; +#[cfg(feature = "client")] +pub mod gateway; +pub mod server; +#[cfg(all(feature = "client", any(test, feature = "testing")))] +pub mod testing; -pub use rmcp::ServiceExt; +pub use rmcp::model; diff --git a/crates/mcp-utils/src/protocol.rs b/crates/mcp-utils/src/protocol.rs deleted file mode 100644 index 578c54f99..000000000 --- a/crates/mcp-utils/src/protocol.rs +++ /dev/null @@ -1,14 +0,0 @@ -use rmcp::{ClientLifecycleMode, model::ProtocolVersion}; - -pub(crate) fn client_lifecycle_mode() -> ClientLifecycleMode { - ClientLifecycleMode::Auto { - preferred_versions: vec![ - ProtocolVersion::V_2026_07_28, - ProtocolVersion::V_2025_11_25, - ProtocolVersion::V_2025_06_18, - ProtocolVersion::V_2025_03_26, - ProtocolVersion::V_2024_11_05, - ], - legacy_version: Some(ProtocolVersion::V_2025_11_25), - } -} diff --git a/crates/mcp-utils/src/server.rs b/crates/mcp-utils/src/server.rs new file mode 100644 index 000000000..cfba4db21 --- /dev/null +++ b/crates/mcp-utils/src/server.rs @@ -0,0 +1,194 @@ +use futures::future::BoxFuture; +use rmcp::service::{DynService, ServerInitializeError, serve_server_with_ct}; +use rmcp::transport::{IntoTransport, stdio}; +use rmcp::{RoleServer, ServerHandler, Service}; +use std::env::{temp_dir, var_os}; +use std::fs::{Permissions, create_dir_all, remove_dir, remove_file, set_permissions}; +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::{fmt, io}; +use thiserror::Error; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::net::UnixListener; +use tokio::task::JoinError; +use tokio_util::sync::{CancellationToken, DropGuard}; +use uuid::Uuid; + +#[derive(Clone)] +pub struct McpServer { + service: Arc Box> + Send + Sync>, +} + +#[derive(Debug)] +pub struct ServerHandle { + connections: CancellationToken, + _stop_on_drop: DropGuard, + socket: SocketPath, +} + +#[derive(Debug)] +pub struct UnixSocket { + listener: UnixListener, + socket: SocketPath, +} + +#[derive(Debug, Error)] +pub enum ServeError { + #[error("Failed to create MCP socket directory {}: {source}", path.display())] + SocketDirectory { + path: PathBuf, + #[source] + source: io::Error, + }, + #[error("Failed to bind MCP socket {}: {source}", path.display())] + Bind { + path: PathBuf, + #[source] + source: io::Error, + }, + #[error("MCP server initialization failed: {0}")] + Init(#[source] Box), + #[error("MCP server task failed: {0}")] + Task(#[source] JoinError), +} + +impl McpServer { + pub fn new(handler: impl ServerHandler) -> Self { + let handler = Arc::new(handler); + Self { service: Arc::new(move || Box::new(Arc::clone(&handler)) as Box>) } + } + + pub async fn serve_stdio(self) -> Result<(), ServeError> { + host(self.service(), stdio(), CancellationToken::new()).await + } + + pub fn serve_unix(self, socket: UnixSocket) -> ServerHandle { + let UnixSocket { listener, socket } = socket; + let connections = CancellationToken::new(); + tokio::spawn(accept_connections(listener, self, connections.clone())); + ServerHandle { _stop_on_drop: connections.clone().drop_guard(), connections, socket } + } + + #[cfg(feature = "client")] + pub(crate) fn serve_in_memory(&self) -> tokio::io::DuplexStream { + let (client, server) = tokio::io::duplex(64 * 1024); + tokio::spawn(serve_connection(self.service(), server, CancellationToken::new())); + client + } + + fn service(&self) -> Box> { + (self.service)() + } +} + +impl From for McpServer { + fn from(handler: H) -> Self { + Self::new(handler) + } +} + +impl fmt::Debug for McpServer { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("McpServer").finish_non_exhaustive() + } +} + +impl ServerHandle { + pub fn path(&self) -> &Path { + self.socket.path() + } + + /// Stops accepting connections, ends the open ones, and removes the socket, as dropping the handle does. + pub fn stop(&self) { + self.connections.cancel(); + self.socket.remove(); + } +} + +impl UnixSocket { + pub fn bind() -> Result { + let socket = SocketPath::allocate()?; + let listener = UnixListener::bind(socket.path()) + .map_err(|source| ServeError::Bind { path: socket.path().to_path_buf(), source })?; + Ok(Self { listener, socket }) + } + + pub fn path(&self) -> &Path { + self.socket.path() + } +} + +#[derive(Debug)] +struct SocketPath { + directory: PathBuf, + socket: PathBuf, +} + +impl SocketPath { + fn allocate() -> Result { + let short_id = &Uuid::new_v4().simple().to_string()[..8]; + let runtime_dir = + var_os("XDG_RUNTIME_DIR").filter(|path| Path::new(path).is_absolute()).map_or_else(temp_dir, PathBuf::from); + let directory = runtime_dir.join("aether").join(format!("aether-{short_id}")); + create_dir_all(&directory).map_err(|source| ServeError::SocketDirectory { path: directory.clone(), source })?; + let path = Self { socket: directory.join("ipc.sock"), directory }; + set_permissions(&path.directory, Permissions::from_mode(0o700)) + .map_err(|source| ServeError::SocketDirectory { path: path.directory.clone(), source })?; + Ok(path) + } + + fn path(&self) -> &Path { + &self.socket + } + + fn remove(&self) { + let _ = remove_file(&self.socket); + let _ = remove_dir(&self.directory); + } +} + +impl Drop for SocketPath { + fn drop(&mut self) { + self.remove(); + } +} + +async fn accept_connections(listener: UnixListener, server: McpServer, cancellation: CancellationToken) { + loop { + let accepted = tokio::select! { + () = cancellation.cancelled() => break, + accepted = listener.accept() => accepted, + }; + let Ok((stream, _)) = accepted else { break }; + tokio::spawn(serve_connection(server.service(), stream, cancellation.child_token())); + } +} + +async fn serve_connection(service: Box>, stream: S, cancellation: CancellationToken) +where + S: AsyncRead + AsyncWrite + Send + 'static, +{ + if let Err(error) = host(service, stream, cancellation).await { + tracing::debug!(%error, "MCP server connection ended with an error"); + } +} + +fn host( + service: T, + transport: U, + cancellation: CancellationToken, +) -> BoxFuture<'static, Result<(), ServeError>> +where + T: Service, + U: IntoTransport, + V: std::error::Error + Send + Sync + 'static, +{ + Box::pin(async move { + let running = serve_server_with_ct(service, transport, cancellation) + .await + .map_err(|error| ServeError::Init(Box::new(error)))?; + running.waiting().await.map_err(ServeError::Task)?; + Ok(()) + }) +} diff --git a/crates/mcp-utils/src/server/mod.rs b/crates/mcp-utils/src/server/mod.rs deleted file mode 100644 index 4c84a4309..000000000 --- a/crates/mcp-utils/src/server/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod mrtr; -pub mod tasks; diff --git a/crates/mcp-utils/src/server/mrtr.rs b/crates/mcp-utils/src/server/mrtr.rs deleted file mode 100644 index 241d17052..000000000 --- a/crates/mcp-utils/src/server/mrtr.rs +++ /dev/null @@ -1,123 +0,0 @@ -use rmcp::model::{ClientCapabilities, ElicitRequestParams, ErrorData, InputRequest, InputRequests, InputResponses}; -use serde::de::DeserializeOwned; - -pub const ELICITATION_UNSUPPORTED: &str = "This tool needs to ask the user for input, but the connected client does not support \ - interactive input (MCP elicitation over protocol 2026-07-28 or newer)."; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ElicitationMode { - Form, - Url, -} - -pub fn input_requests_supported(capabilities: Option<&ClientCapabilities>, requests: &InputRequests) -> bool { - requests.values().all(|request| capabilities.is_some_and(|capabilities| supports_input(capabilities, request))) -} - -pub fn elicitation_supported(capabilities: Option<&ClientCapabilities>, mode: ElicitationMode) -> bool { - let Some(elicitation) = capabilities.and_then(|capabilities| capabilities.elicitation.as_ref()) else { - return false; - }; - match mode { - ElicitationMode::Form => elicitation.form.is_some() || elicitation.url.is_none(), - ElicitationMode::Url => elicitation.url.is_some(), - } -} - -/// Deserialize one keyed response from a complete MRTR response batch. -pub fn parse_response(responses: &InputResponses, key: &str) -> Result { - let response = responses - .get(key) - .ok_or_else(|| ErrorData::invalid_params(format!("missing input response for '{key}'"), None))?; - serde_json::from_value(response.clone()) - .map_err(|e| ErrorData::invalid_params(format!("invalid input response for '{key}': {e}"), None)) -} - -fn supports_input(capabilities: &ClientCapabilities, request: &InputRequest) -> bool { - #[allow(unreachable_patterns)] - match request { - InputRequest::CreateMessage(_) => capabilities.sampling.is_some(), - InputRequest::Elicitation(request) => supports_elicitation(capabilities, &request.params), - InputRequest::ListRoots(_) => capabilities.roots.is_some(), - _ => false, - } -} - -fn supports_elicitation(capabilities: &ClientCapabilities, params: &ElicitRequestParams) -> bool { - #[allow(unreachable_patterns)] - let mode = match params { - ElicitRequestParams::FormElicitationParams { .. } => ElicitationMode::Form, - ElicitRequestParams::UrlElicitationParams { .. } => ElicitationMode::Url, - _ => return false, - }; - elicitation_supported(Some(capabilities), mode) -} - -#[cfg(test)] -mod tests { - use super::*; - use rmcp::model::{ - ClientCapabilities, ElicitRequest, ElicitRequestParams, ElicitationCapability, ElicitationSchema, - FormElicitationCapability, InputRequest, UrlElicitationCapability, - }; - use serde_json::json; - - #[test] - fn validates_capabilities_against_every_request_in_the_batch() { - let requests = InputRequests::from([ - ("form".to_string(), elicitation("Form")), - ("url".to_string(), url_elicitation("URL")), - ]); - let mut capabilities = ClientCapabilities::default(); - capabilities.elicitation = Some(ElicitationCapability::new().with_form(FormElicitationCapability::new())); - - assert!(!input_requests_supported(Some(&capabilities), &requests)); - } - - #[test] - fn validates_a_mixed_elicitation_batch_when_all_capabilities_are_present() { - let requests = InputRequests::from([ - ("form".to_string(), elicitation("Form")), - ("url".to_string(), url_elicitation("URL")), - ]); - let mut capabilities = ClientCapabilities::default(); - capabilities.elicitation = Some( - ElicitationCapability::new() - .with_form(FormElicitationCapability::new()) - .with_url(UrlElicitationCapability::new()), - ); - - assert!(input_requests_supported(Some(&capabilities), &requests)); - } - - #[test] - fn parse_response_deserializes_a_key_from_the_batch() { - let responses = InputResponses::from([("answer".to_string(), json!(42))]); - - assert_eq!(parse_response::(&responses, "answer").unwrap(), 42); - } - - #[test] - fn parse_response_rejects_a_missing_key() { - let error = parse_response::(&InputResponses::new(), "answer").unwrap_err(); - - assert_eq!(error.message, "missing input response for 'answer'"); - } - - fn elicitation(message: &str) -> InputRequest { - InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::FormElicitationParams { - meta: None, - message: message.to_string(), - requested_schema: ElicitationSchema::builder().build().unwrap(), - })) - } - - fn url_elicitation(message: &str) -> InputRequest { - InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::UrlElicitationParams { - meta: None, - message: message.to_string(), - url: "https://example.com/input".to_string(), - elicitation_id: "id".to_string(), - })) - } -} diff --git a/crates/mcp-utils/src/server/tasks.rs b/crates/mcp-utils/src/server/tasks.rs deleted file mode 100644 index c6d4cdf0f..000000000 --- a/crates/mcp-utils/src/server/tasks.rs +++ /dev/null @@ -1,15 +0,0 @@ -use rmcp::model::ClientCapabilities; -use rmcp::service::RequestContext; -use rmcp::{ErrorData, RoleServer}; - -/// Bounds both task execution and terminal-result retention because rmcp uses -/// one TTL for both concerns. -pub const BACKGROUND_TASK_TTL_MS: u64 = 3_600_000; - -pub fn require_tasks_capability(context: &RequestContext) -> Result<(), ErrorData> { - if context.client_capabilities().is_some_and(|capabilities| capabilities.supports_tasks()) { - Ok(()) - } else { - Err(ErrorData::missing_required_client_capability(ClientCapabilities::builder().enable_tasks().build())) - } -} diff --git a/crates/mcp-utils/src/server/tool.rs b/crates/mcp-utils/src/server/tool.rs deleted file mode 100644 index e69de29bb..000000000 diff --git a/crates/mcp-utils/src/testing.rs b/crates/mcp-utils/src/testing.rs index 4e4ae48b2..77375a841 100644 --- a/crates/mcp-utils/src/testing.rs +++ b/crates/mcp-utils/src/testing.rs @@ -1,242 +1,50 @@ -use std::future::Future; - -use crate::protocol::client_lifecycle_mode; -use rmcp::{ - RoleClient, RoleServer, Service, serve_client_with_lifecycle, serve_server, - service::{ClientInitializeError, RunningService, ServerInitializeError}, +pub use elicitation_script::{ + CapturedElicitation, ElicitationScript, ElicitationScriptBuilder, elicitation_host, elicitation_sink, +}; +pub use fake_auth_server::{ + ACCESS_TOKEN, CLIENT_METADATA_URL, FakeAuthServer, OAuthPrompt, client_metadata, complete_authorization, }; - -#[cfg(feature = "client")] -pub use elicitation_script::{CapturedElicitation, ElicitationScript, UrlElicitationHandler, url_elicitation_handler}; -#[cfg(all(feature = "client", any(test, feature = "testing")))] pub use fake_mcp::{ CapturedTaskUpdate, CapturedToolCall, FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, - completed_task_payload, fake_mcp, + completed_task_payload, }; +pub use fake_task::{ + FakeTask, accept_with_name, cancelled_task, completed_task, failed_task, form_request, input_requests, + input_required_task, name_input, task, url_request, working_task, +}; +pub use gate::Gate; -#[cfg(all(feature = "client", any(test, feature = "testing")))] +mod elicitation_script; +mod fake_auth_server; mod fake_mcp; - -pub type ConnectedServices = (RunningService, RunningService); - -/// Helper function to connect an MCP server and client via in-memory transport -/// This handles the dual-era discovery/initialization handshake by running both concurrently -pub fn connect(server: T, client: U) -> impl Future, ConnectError>> -where - T: Service, - U: Service, -{ - Box::pin(async move { - let (client_transport, server_transport) = tokio::io::duplex(64 * 1024); - - let server = tokio::spawn(serve_server(server, server_transport)); - let client = serve_client_with_lifecycle(client, client_transport, client_lifecycle_mode()) - .await - .map_err(|error| ConnectError::ClientInit(Box::new(error)))?; - - let _ = client.list_tools(None).await; - let server = server.await?.map_err(|error| ConnectError::ServerInit(Box::new(error)))?; - - Ok((server, client)) - }) +mod fake_task; +mod gate; + +use crate::client::{ClientOptions, McpClient, ToolCallError, ToolCallOptions, Transport}; +use crate::gateway::McpGateway; +use crate::model::CallToolResult; +use crate::server::McpServer; +use futures::StreamExt; +use serde_json::{Map, Value}; + +/// Connects an [`McpClient`] named `name` to an in-process `server`. +pub async fn connect(name: &str, server: impl Into, options: &ClientOptions) -> McpClient { + McpClient::connect(name, Transport::InProcess(server.into()), options).await.expect("connect in-process server") } -#[derive(Debug, thiserror::Error)] -pub enum ConnectError { - #[error("Server initialization failed: {0}")] - ServerInit(Box), - #[error("Client initialization failed: {0}")] - ClientInit(Box), - #[error("Server task failed: {0}")] - ServerTask(#[from] tokio::task::JoinError), +/// Calls `name` with default options and waits for its result. +pub async fn call(client: &McpClient, name: &str, arguments: Value) -> Result { + client.call_tool(name, args(arguments), ToolCallOptions::default()).result().await } -#[cfg(feature = "client")] -mod elicitation_script { - use crate::client::McpClientEvent; - use futures::future::BoxFuture; - use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction}; - use std::collections::VecDeque; - use std::future::Future; - use std::sync::{Arc, Mutex, PoisonError}; - use tokio::sync::mpsc; - use tokio::task::JoinHandle; - - pub type UrlElicitationHandler = Arc BoxFuture<'static, ()> + Send + Sync>; - - /// Scripts the user's side of elicitation round trips: answers each - /// incoming request with the next queued response (Cancel once the queue - /// is empty) and records what arrived for assertions. - pub struct ElicitationScript { - captured: Arc>>, - task: JoinHandle<()>, - } - - pub fn url_elicitation_handler(handler: F) -> UrlElicitationHandler - where - F: Fn(String, String) -> Fut + Send + Sync + 'static, - Fut: Future + Send + 'static, - { - Arc::new(move |url, id| Box::pin(handler(url, id))) - } - - #[derive(Clone)] - pub struct CapturedElicitation { - pub server_name: String, - pub request: ElicitRequestParams, - } - - impl ElicitationScript { - pub fn spawn( - event_rx: mpsc::Receiver, - responses: impl IntoIterator, - ) -> Self { - Self::spawn_with_url_handler(event_rx, responses, None) - } - - pub fn spawn_with_url_handler( - mut event_rx: mpsc::Receiver, - responses: impl IntoIterator, - on_url: Option, - ) -> Self { - let mut responses = responses.into_iter().collect::>(); - let captured = Arc::new(Mutex::new(Vec::new())); - let recorder = Arc::clone(&captured); - let task = tokio::spawn(async move { - while let Some(event) = event_rx.recv().await { - let McpClientEvent::Elicitation(event) = event else { continue }; - recorder - .lock() - .unwrap_or_else(PoisonError::into_inner) - .push(CapturedElicitation { server_name: event.server_name, request: event.request.clone() }); - if let (Some(on_url), ElicitRequestParams::UrlElicitationParams { url, elicitation_id, .. }) = - (&on_url, event.request) - { - on_url(url, elicitation_id).await; - } - let response = - responses.pop_front().unwrap_or_else(|| ElicitResult::new(ElicitationAction::Cancel)); - let _ = event.response_sender.send(response); - } - }); - Self { captured, task } - } - - pub fn captured(&self) -> Vec { - self.captured.lock().unwrap_or_else(PoisonError::into_inner).clone() - } - } - - impl Drop for ElicitationScript { - fn drop(&mut self) { - self.task.abort(); - } +pub async fn wait_until(gateway: &McpGateway, condition: impl Fn(&McpGateway) -> bool) { + let mut changes = gateway.changes(); + while !condition(gateway) { + changes.next().await.expect("the gateway stopped before the condition held"); } } -#[cfg(test)] -mod tests { - use super::connect; - use rmcp::{ - ClientHandler, ServerHandler, - model::{ - ErrorData, Implementation, InitializeRequestParams, ProtocolVersion, ServerCapabilities, ServerConfig, - }, - service::RequestContext, - }; - use std::borrow::Cow; - - #[tokio::test] - async fn connect_prefers_stateless_discovery_for_modern_servers() { - let (_server, client) = connect(McpServer728, TestClient).await.expect("connect"); - assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2026_07_28); - client.list_tools(None).await.expect("list tools"); - client.cancel().await.expect("cancel client"); - } - - #[tokio::test] - async fn connect_selects_an_older_mutually_supported_revision() { - let (_server, client) = connect(McpServer618, TestClient).await.expect("connect"); - - assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2025_06_18); - client.list_tools(None).await.expect("list tools"); - client.cancel().await.expect("cancel client"); - } - - #[tokio::test] - async fn connect_falls_back_to_legacy_initialization() { - let (_server, client) = connect(McpServer1125, TestClient).await.expect("connect"); - - assert_eq!(client.peer_info().expect("peer info").protocol_version, ProtocolVersion::V_2025_11_25); - client.cancel().await.expect("cancel client"); - } - - #[derive(Clone, Default)] - struct TestClient; - - impl ClientHandler for TestClient {} - - #[derive(Clone, Default)] - struct McpServer728; - - impl ServerHandler for McpServer728 { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("modern-only", "1.0.0")) - .with_protocol_version(ProtocolVersion::V_2026_07_28) - } - - fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { - Cow::Owned(vec![ProtocolVersion::V_2026_07_28]) - } - - fn initialize( - &self, - _request: InitializeRequestParams, - _context: RequestContext, - ) -> impl std::future::Future> + Send + '_ { - std::future::ready(Err(ErrorData::new( - rmcp::model::ErrorCode::METHOD_NOT_FOUND, - "initialize is not supported", - None, - ))) - } - } - - #[derive(Clone, Default)] - struct McpServer618; - - impl ServerHandler for McpServer618 { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("older-revision", "1.0.0")) - .with_protocol_version(ProtocolVersion::V_2025_06_18) - } - - fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { - Cow::Owned(vec![ProtocolVersion::V_2025_06_18]) - } - } - - #[derive(Clone, Default)] - struct McpServer1125; - - impl ServerHandler for McpServer1125 { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - .with_server_info(Implementation::new("legacy", "1.0.0")) - .with_protocol_version(ProtocolVersion::V_2025_11_25) - } - - fn discover( - &self, - _context: RequestContext, - ) -> impl Future> + Send + '_ { - std::future::ready(Err(ErrorData::new( - rmcp::model::ErrorCode::METHOD_NOT_FOUND, - "server/discover is not supported", - None, - ))) - } - } +pub fn args(value: Value) -> Map { + let Value::Object(arguments) = value else { panic!("tool arguments must be a JSON object") }; + arguments } diff --git a/crates/mcp-utils/src/testing/elicitation_script.rs b/crates/mcp-utils/src/testing/elicitation_script.rs new file mode 100644 index 000000000..0f9776710 --- /dev/null +++ b/crates/mcp-utils/src/testing/elicitation_script.rs @@ -0,0 +1,106 @@ +use crate::client::{ClientOptions, Elicitation, ElicitationSink, cancelled}; +use futures::future::BoxFuture; +use rmcp::model::{ElicitRequestParams, ElicitResult}; +use std::collections::VecDeque; +use std::future::Future; +use std::sync::{Arc, Mutex, PoisonError}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +pub struct ElicitationScript { + captured: Arc>>, + task: JoinHandle<()>, +} + +#[derive(Default)] +pub struct ElicitationScriptBuilder { + responses: VecDeque, + on_url: Option, +} + +#[derive(Clone)] +pub struct CapturedElicitation { + pub server_name: String, + pub request: ElicitRequestParams, +} + +impl ElicitationScript { + pub fn builder() -> ElicitationScriptBuilder { + ElicitationScriptBuilder::default() + } + + pub fn answering(responses: impl IntoIterator) -> (ClientOptions, Self) { + let (sink, script) = Self::builder().responses(responses).sink(); + (ClientOptions::default().elicitation(sink), script) + } + + pub fn captured(&self) -> Vec { + self.captured.lock().unwrap_or_else(PoisonError::into_inner).clone() + } +} + +impl Drop for ElicitationScript { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl ElicitationScriptBuilder { + pub fn response(mut self, response: ElicitResult) -> Self { + self.responses.push_back(response); + self + } + + pub fn responses(mut self, responses: impl IntoIterator) -> Self { + self.responses.extend(responses); + self + } + + pub fn on_url(mut self, handler: F) -> Self + where + F: Fn(String, String) -> Fut + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + self.on_url = Some(Arc::new(move |url, id| Box::pin(handler(url, id)))); + self + } + + pub fn sink(self) -> (ElicitationSink, ElicitationScript) { + let (sink, events) = elicitation_sink(); + (sink, self.spawn(events)) + } + + pub fn spawn(self, mut event_rx: mpsc::Receiver) -> ElicitationScript { + let Self { mut responses, on_url } = self; + let captured = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::clone(&captured); + let task = tokio::spawn(async move { + while let Some(event) = event_rx.recv().await { + let Elicitation::Request(elicitation) = event else { continue }; + recorder.lock().unwrap_or_else(PoisonError::into_inner).push(CapturedElicitation { + server_name: elicitation.server.clone(), + request: elicitation.request.clone(), + }); + if let (Some(on_url), ElicitRequestParams::UrlElicitationParams { url, elicitation_id, .. }) = + (&on_url, &elicitation.request) + { + on_url(url.clone(), elicitation_id.clone()).await; + } + elicitation.respond(responses.pop_front().unwrap_or_else(cancelled)); + } + }); + ElicitationScript { captured, task } + } +} + +pub fn elicitation_sink() -> (ElicitationSink, mpsc::Receiver) { + let (events, host) = mpsc::channel(32); + (ElicitationSink::form_and_url(events), host) +} + +pub fn elicitation_host() -> (ClientOptions, mpsc::Receiver) { + let (sink, host) = elicitation_sink(); + (ClientOptions::default().elicitation(sink), host) +} + +type UrlHandler = Arc BoxFuture<'static, ()> + Send + Sync>; diff --git a/crates/mcp-utils/src/testing/fake_auth_server.rs b/crates/mcp-utils/src/testing/fake_auth_server.rs new file mode 100644 index 000000000..4bdd80ea2 --- /dev/null +++ b/crates/mcp-utils/src/testing/fake_auth_server.rs @@ -0,0 +1,172 @@ +use super::FakeMcpServer; +use crate::client::{Elicitation, ElicitationRequest, OAuthSettings, Transport}; +use crate::gateway::ServerSpec; +use aether_auth::OAuthClientRegistration; +use axum::extract::{Request, State}; +use axum::http::{StatusCode, Uri, header::WWW_AUTHENTICATE}; +use axum::middleware::{Next, from_fn_with_state}; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use reqwest::Url; +use reqwest::header::{AUTHORIZATION, HeaderMap}; +use rmcp::model::{ElicitRequestParams, ElicitResult, ElicitationAction}; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; +use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; +use tokio::net::TcpListener; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; + +pub struct FakeAuthServer { + url: String, + auth: Arc, + task: JoinHandle<()>, +} + +pub struct OAuthPrompt { + pub elicitation: Box, + pub url: Url, +} + +pub const CLIENT_METADATA_URL: &str = "https://client.example/oauth/client-metadata.json"; + +impl FakeAuthServer { + pub async fn start(server: FakeMcpServer) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind the fake authorization server"); + let origin = format!("http://{}", listener.local_addr().expect("a bound address")); + let auth = Arc::new(AuthState { origin: origin.clone(), bearers: Mutex::new(Vec::new()) }); + let mcp = StreamableHttpService::new( + move || Ok(server.clone()), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + let protected = + Router::new().nest_service("/mcp", mcp).layer(from_fn_with_state(Arc::clone(&auth), require_token)); + let app = Router::new() + .route("/token", post(token)) + .fallback(metadata) + .with_state(Arc::clone(&auth)) + .merge(protected); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve the fake authorization server"); + }); + Self { url: format!("{origin}/mcp"), auth, task } + } + + pub fn transport(&self) -> Transport { + self.transport_with(Some(client_metadata())) + } + + pub fn transport_with(&self, oauth: Option) -> Transport { + Transport::Http { url: self.url.clone(), headers: HeaderMap::new(), oauth } + } + + pub fn spec(&self, name: &str) -> ServerSpec { + ServerSpec { name: name.to_string(), transport: self.transport() } + } + + pub fn received_bearer(&self, token: &str) -> bool { + self.auth.bearers.lock().expect("bearers lock").iter().any(|bearer| bearer == token) + } +} + +impl Drop for FakeAuthServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl OAuthPrompt { + pub async fn receive(host: &mut mpsc::Receiver) -> Self { + let Some(Elicitation::Request(elicitation)) = host.recv().await else { + panic!("expected the authorization prompt") + }; + let ElicitRequestParams::UrlElicitationParams { url, elicitation_id, .. } = &elicitation.request else { + panic!("expected a URL elicitation") + }; + assert_eq!(elicitation_id, "oauth"); + let url = Url::parse(url).expect("an authorization URL"); + Self { elicitation, url } + } + + pub fn query(&self, name: &str) -> String { + self.url.query_pairs().find(|(key, _)| key == name).map(|(_, value)| value.into_owned()).unwrap_or_default() + } + + pub fn callback_url(&self) -> String { + callback_url(self.url.as_str()) + } + + pub async fn approve(self) { + let url = self.url.to_string(); + self.elicitation.respond(ElicitResult::new(ElicitationAction::Accept)); + complete_authorization(&url).await; + } +} + +pub async fn complete_authorization(authorization_url: &str) { + reqwest::get(callback_url(authorization_url)).await.expect("the loopback callback responds"); +} + +fn callback_url(authorization_url: &str) -> String { + let url = Url::parse(authorization_url).expect("an authorization URL"); + let query = |name: &str| url.query_pairs().find(|(key, _)| key == name).map(|(_, value)| value.into_owned()); + let mut callback = Url::parse(&query("redirect_uri").expect("a redirect URI")).expect("a valid redirect URI"); + callback + .query_pairs_mut() + .append_pair("code", "test-code") + .append_pair("state", &query("state").unwrap_or_default()); + callback.to_string() +} + +pub fn client_metadata() -> OAuthSettings { + OAuthSettings { + registration: OAuthClientRegistration::ClientMetadata(CLIENT_METADATA_URL.to_string()), + callback_port: 0, + } +} + +pub const ACCESS_TOKEN: &str = "access-token"; + +struct AuthState { + origin: String, + bearers: Mutex>, +} + +async fn require_token(State(auth): State>, request: Request, next: Next) -> Response { + let bearer = request + .headers() + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::to_string); + if let Some(bearer) = bearer { + auth.bearers.lock().expect("bearers lock").push(bearer.clone()); + if bearer == ACCESS_TOKEN { + return next.run(request).await; + } + } + let challenge = format!("Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\"", auth.origin); + (StatusCode::UNAUTHORIZED, [(WWW_AUTHENTICATE, challenge)]).into_response() +} + +async fn token() -> Json { + Json(json!({ "access_token": ACCESS_TOKEN, "token_type": "Bearer", "expires_in": 3600 })) +} + +async fn metadata(State(auth): State>, uri: Uri) -> Json { + let origin = &auth.origin; + if uri.path().contains("oauth-protected-resource") { + return Json(json!({ "resource": format!("{origin}/mcp"), "authorization_servers": [origin] })); + } + Json(json!({ + "issuer": origin, + "authorization_endpoint": format!("{origin}/authorize"), + "token_endpoint": format!("{origin}/token"), + "response_types_supported": ["code"], + "code_challenge_methods_supported": ["S256"], + "client_id_metadata_document_supported": true + })) +} diff --git a/crates/mcp-utils/src/testing/fake_mcp.rs b/crates/mcp-utils/src/testing/fake_mcp.rs index b19f424dd..d05d015e0 100644 --- a/crates/mcp-utils/src/testing/fake_mcp.rs +++ b/crates/mcp-utils/src/testing/fake_mcp.rs @@ -1,45 +1,46 @@ -use crate::client::{RuntimeMcpServer, RuntimeMcpTransport, ToolExposure}; use rmcp::{ ErrorData as McpError, Peer, RoleServer, ServerHandler, model::{ CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ClientCapabilities, - ContentBlock, CreateTaskResult, DetailedTask, DiscoverResult, GetTaskParams, GetTaskResult, Implementation, - ListToolsResult, PaginatedRequestParams, ProgressNotificationParam, ProtocolVersion, ResultType, - ServerCapabilities, ServerConfig, TaskPayload, Tool, UpdateTaskParams, + ContentBlock, CreateTaskResult, DetailedTask, DiscoverResult, ElicitResult, GetPromptRequestParams, + GetPromptResponse, GetPromptResult, GetTaskParams, GetTaskResult, Implementation, ListPromptsResult, + ListToolsResult, PaginatedRequestParams, ProgressNotificationParam, ProgressToken, Prompt, PromptMessage, + PromptsCapability, ProtocolVersion, ResultType, Role, ServerCapabilities, ServerConfig, TaskPayload, Tool, + ToolAnnotations, UpdateTaskParams, }, - service::{DynService, RequestContext}, + service::RequestContext, }; -use serde_json::json; +use serde_json::{Value, json}; use std::collections::{BTreeMap, HashMap, VecDeque}; use std::future::Future; -use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::sync::{Arc, Mutex, Weak}; +use tokio::sync::{mpsc, watch}; +use tokio::time::Instant; -pub fn fake_mcp(name: &str, server: FakeMcpServer) -> RuntimeMcpServer { - RuntimeMcpServer::new(name, RuntimeMcpTransport::InMemory { server: server.into_dyn() }, ToolExposure::ModelVisible) -} +use super::Gate; pub fn completed_task_payload(result: CallToolResult) -> TaskPayload { let result = serde_json::to_value(result).and_then(serde_json::from_value).expect("a tool result is a JSON object"); TaskPayload::Completed { result } } -/// A fake MCP server preloaded with the classic math tools (`add_numbers`, -/// `divide_numbers`, `slow_tool`); add scripted tools with [`Self::with_tool`]. #[derive(Clone)] pub struct FakeMcpServer { state: FakeMcpState, + alive: Arc<()>, } #[derive(Clone, Default)] pub struct FakeMcpState { inner: Arc>, + changes: Arc>, } #[derive(Clone)] pub struct CapturedToolCall { pub request: CallToolRequestParams, pub context_meta: serde_json::Map, + pub at: Instant, } #[derive(Clone)] @@ -58,21 +59,22 @@ pub struct FakeTool { #[derive(Clone)] pub struct FakeToolResponse { response: CallToolResponse, - delay: Duration, + gate: Option, progress: Vec, - task_progress: Vec<(f64, Option)>, -} - -#[derive(Clone)] -struct FakeProgress { - progress: f64, - total: Option, - message: Option, } impl FakeMcpServer { pub fn new() -> Self { - Self::default() + Self { state: FakeMcpState::default(), alive: Arc::default() } + } + + pub fn math() -> Self { + Self::new().with_tool(add_numbers()).with_tool(divide_numbers()) + } + + pub fn with_instructions(self, instructions: impl Into) -> Self { + self.state.lock().instructions = Some(instructions.into()); + self } pub fn with_tool(self, tool: FakeTool) -> Self { @@ -80,6 +82,31 @@ impl FakeMcpServer { self } + pub fn with_prompt(self, name: impl Into, template: impl Into) -> Self { + let name = name.into(); + self.with_prompt_definition(Prompt::new(name, Some("Fake prompt"), None), template) + } + + pub fn with_prompt_definition(self, prompt: Prompt, template: impl Into) -> Self { + self.state.lock().prompts.insert(prompt.name.clone(), (prompt, template.into())); + self + } + + pub fn with_prompt_gate(self, gate: Gate) -> Self { + self.state.lock().prompt_gate = Some(gate); + self + } + + pub fn with_tool_list_gate(self, gate: Gate) -> Self { + self.state.lock().tool_list_gate = Some(gate); + self + } + + pub fn with_prompt_elicitation(self, results: mpsc::UnboundedSender) -> Self { + self.state.lock().prompt_elicitations = Some(results); + self + } + pub fn with_task(self, task_id: impl Into, states: impl IntoIterator) -> Self { self.state.script_task(task_id, states); self @@ -99,8 +126,8 @@ impl FakeMcpServer { self.state.clone() } - pub fn into_dyn(self) -> Box> { - Box::new(self) + pub fn alive(&self) -> Weak<()> { + Arc::downgrade(&self.alive) } } @@ -109,6 +136,10 @@ impl FakeMcpState { self.lock().calls.iter().filter(|call| call.request.name.as_ref() == tool).cloned().collect() } + pub fn tool_list_requests(&self) -> usize { + self.lock().tool_list_requests + } + pub fn task_get_ids(&self) -> Vec { self.lock().task_get_ids.clone() } @@ -129,64 +160,59 @@ impl FakeMcpState { self.lock().tasks.insert(task_id.into(), states.into_iter().collect()); } - fn task_for(&self, task_id: &str) -> Result, ()> { - let mut inner = self.lock(); - inner.task_get_ids.push(task_id.to_string()); - if inner.task_get_failures > 0 { - inner.task_get_failures -= 1; - return Err(()); - } - let Some(states) = inner.tasks.get_mut(task_id) else { - return Ok(None); - }; - Ok(if states.len() > 1 { states.pop_front() } else { states.front().cloned() }) + pub fn add_tool(&self, tool: FakeTool) { + self.lock().tools.insert(tool.definition.name.to_string(), tool); } - fn record_task_update(&self, request: UpdateTaskParams) -> bool { - let mut inner = self.lock(); - inner - .task_updates - .push(CapturedTaskUpdate { task_id: request.task_id, input_responses: request.input_responses }); - if inner.task_update_failures > 0 { - inner.task_update_failures -= 1; - false - } else { - true - } + pub async fn add_tool_and_notify(&self, tool: FakeTool) { + self.mutate_tools_and_notify(|tools| { + tools.insert(tool.definition.name.to_string(), tool); + }) + .await; } - fn record_task_cancel(&self, request: CancelTaskParams) { - self.lock().task_cancel_ids.push(request.task_id); + pub async fn clear_tools_and_notify(&self) { + self.mutate_tools_and_notify(BTreeMap::clear).await; } - pub fn add_tool(&self, tool: FakeTool) { - self.lock().tools.insert(tool.definition.name.to_string(), tool); + pub async fn notify_tools_changed(&self) { + self.mutate_tools_and_notify(|_| {}).await; } - pub async fn add_tool_and_notify(&self, tool: FakeTool) { - let peers = { - let mut inner = self.lock(); - inner.tools.insert(tool.definition.name.to_string(), tool); - inner.peers.clone() - }; - for peer in peers { - let _ = peer.notify_tool_list_changed().await; + pub fn fail_next_tool_list(&self) { + self.lock().tool_list_failures += 1; + } + + pub async fn wait_until(&self, condition: impl Fn(&Self) -> bool) { + let mut changes = self.changes.subscribe(); + while !condition(self) { + changes.changed().await.expect("the state outlives its waiters"); } } - pub async fn clear_tools_and_notify(&self) { - let peers = { + async fn mutate_tools_and_notify(&self, mutate: impl FnOnce(&mut BTreeMap)) { + let client = { let mut inner = self.lock(); - inner.tools.clear(); - inner.peers.clone() + mutate(&mut inner.tools); + inner.client.clone() }; - for peer in peers { - let _ = peer.notify_tool_list_changed().await; + if let Some(client) = client { + let _ = client.notify_tool_list_changed().await; } } - pub fn fail_next_tool_list(&self) { - self.lock().tool_list_failures += 1; + fn list_tools(&self, peer: Peer) -> Result, McpError> { + self.record(|inner| { + inner.tool_list_requests += 1; + if inner.tool_list_failures > 0 { + inner.tool_list_failures -= 1; + return Err(McpError::internal_error("scripted tools/list failure", None)); + } + if inner.client.as_ref().is_none_or(Peer::is_transport_closed) { + inner.client = Some(peer); + } + Ok(inner.tools.values().map(|tool| tool.definition.clone()).collect()) + }) } fn definitions(&self) -> Vec { @@ -198,14 +224,53 @@ impl FakeMcpState { request: &CallToolRequestParams, context_meta: serde_json::Map, ) -> Option { - let mut inner = self.lock(); - inner.calls.push(CapturedToolCall { request: request.clone(), context_meta }); - inner.tools.get(request.name.as_ref()).and_then(|tool| tool.response_for(request)) + self.record(|inner| { + inner.calls.push(CapturedToolCall { request: request.clone(), context_meta, at: Instant::now() }); + inner.tools.get(request.name.as_ref()).and_then(|tool| tool.response_for(request)) + }) + } + + fn task_for(&self, task_id: &str) -> Result, ()> { + self.record(|inner| { + inner.task_get_ids.push(task_id.to_string()); + if inner.task_get_failures > 0 { + inner.task_get_failures -= 1; + return Err(()); + } + let Some(states) = inner.tasks.get_mut(task_id) else { + return Ok(None); + }; + Ok(if states.len() > 1 { states.pop_front() } else { states.front().cloned() }) + }) + } + + fn record_task_update(&self, request: UpdateTaskParams) -> bool { + self.record(|inner| { + inner + .task_updates + .push(CapturedTaskUpdate { task_id: request.task_id, input_responses: request.input_responses }); + if inner.task_update_failures > 0 { + inner.task_update_failures -= 1; + false + } else { + true + } + }) + } + + fn record_task_cancel(&self, request: CancelTaskParams) { + self.record(|inner| inner.task_cancel_ids.push(request.task_id)); } fn lock(&self) -> std::sync::MutexGuard<'_, FakeMcpStateInner> { self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner) } + + fn record(&self, change: impl FnOnce(&mut FakeMcpStateInner) -> T) -> T { + let result = change(&mut self.lock()); + self.changes.send_replace(()); + result + } } impl FakeTool { @@ -225,6 +290,17 @@ impl FakeTool { self } + pub fn input_schema(mut self, schema: Value) -> Self { + let Value::Object(schema) = schema else { panic!("a tool input schema must be a JSON object") }; + self.definition.input_schema = Arc::new(schema); + self + } + + pub fn annotations(mut self, annotations: ToolAnnotations) -> Self { + self.definition.annotations = Some(annotations); + self + } + pub fn responds(mut self, response: impl Into) -> Self { self.responses.insert(None, response.into()); self @@ -255,7 +331,11 @@ impl FakeTool { impl FakeToolResponse { pub fn new(response: impl Into) -> Self { - Self { response: response.into(), delay: Duration::ZERO, progress: Vec::new(), task_progress: Vec::new() } + Self { response: response.into(), gate: None, progress: Vec::new() } + } + + pub fn pending() -> Self { + Self::text("never sent").gated(Gate::new()) } pub fn text(text: impl Into) -> Self { @@ -266,8 +346,8 @@ impl FakeToolResponse { Self::new(CallToolResponse::Task(task)) } - pub fn delay(mut self, delay: Duration) -> Self { - self.delay = delay; + pub fn gated(mut self, gate: Gate) -> Self { + self.gate = Some(gate); self } @@ -280,11 +360,6 @@ impl FakeToolResponse { self.progress.push(FakeProgress { progress, total: None, message: Some(message.into()) }); self } - - pub fn task_progress(mut self, progress: f64, total: Option) -> Self { - self.task_progress.push((progress, total)); - self - } } impl From for FakeToolResponse @@ -298,10 +373,7 @@ where impl Default for FakeMcpServer { fn default() -> Self { - Self { state: FakeMcpState::default() } - .with_tool(add_numbers()) - .with_tool(divide_numbers()) - .with_tool(slow_tool()) + Self::new() } } @@ -318,11 +390,18 @@ impl ServerHandler for FakeMcpServer { } fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().enable_tasks().build()) - .with_server_info( - Implementation::new("fake-mcp-server", "0.1.0").with_description("A fake MCP server for testing"), - ) - .with_instructions("A fake MCP server for testing") + let inner = self.state.lock(); + let mut capabilities = ServerCapabilities::builder().enable_tools().enable_tasks().build(); + if !inner.prompts.is_empty() { + capabilities.prompts = Some(PromptsCapability::default()); + } + let config = ServerConfig::new(capabilities).with_server_info( + Implementation::new("fake-mcp-server", "0.1.0").with_description("A fake MCP server for testing"), + ); + match &inner.instructions { + Some(instructions) => config.with_instructions(instructions.clone()), + None => config, + } } fn get_task( @@ -360,32 +439,65 @@ impl ServerHandler for FakeMcpServer { std::future::ready(Ok(())) } - fn list_tools( + async fn list_tools( &self, _request: Option, context: RequestContext, - ) -> impl Future> + Send + '_ { + ) -> Result { + let gate = self.state.lock().tool_list_gate.clone(); + if let Some(gate) = gate { + gate.pass().await; + } let supports_cache_hints = context.protocol_version().is_some_and(|version| version >= ProtocolVersion::V_2026_07_28); - let tools = { - let mut inner = self.state.lock(); - if inner.tool_list_failures > 0 { - inner.tool_list_failures -= 1; - return std::future::ready(Err(McpError::internal_error("scripted tools/list failure", None))); - } - if inner.peers.is_empty() { - inner.peers.push(context.peer); - } - inner.tools.values().map(|tool| tool.definition.clone()).collect() - }; - std::future::ready(Ok(ListToolsResult { + self.state.list_tools(context.peer).map(|tools| ListToolsResult { result_type: Some(ResultType::COMPLETE), tools, meta: None, next_cursor: None, ttl_ms: supports_cache_hints.then_some(0), cache_scope: supports_cache_hints.then_some(CacheScope::Public), - })) + }) + } + + fn list_prompts( + &self, + _request: Option, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + let prompts = self.state.lock().prompts.values().map(|(prompt, _)| prompt.clone()).collect(); + std::future::ready(Ok(ListPromptsResult::with_all_items(prompts))) + } + + async fn get_prompt( + &self, + request: GetPromptRequestParams, + context: RequestContext, + ) -> Result { + let (template, gate, elicitations) = { + let inner = self.state.lock(); + let template = inner.prompts.get(&request.name).map(|(_, template)| template.clone()); + (template, inner.prompt_gate.clone(), inner.prompt_elicitations.clone()) + }; + if let Some(gate) = gate { + gate.pass().await; + } + if let Some(results) = elicitations { + let params = serde_json::from_value(json!({ + "mode": "form", "message": "Continue?", "requestedSchema": {"type": "object", "properties": {}} + })) + .expect("valid elicitation request"); + let response = context.peer.create_elicitation(params).await.expect("elicitation response"); + let _ = results.send(response); + } + template + .map(|template| { + let text = request.arguments.unwrap_or_default().iter().fold(template, |text, (name, value)| { + text.replace(&format!("{{{name}}}"), value.as_str().unwrap_or_default()) + }); + GetPromptResult::new(vec![PromptMessage::new_text(Role::User, text)]).into() + }) + .ok_or_else(|| McpError::invalid_params(format!("unknown prompt: {}", request.name), None)) } fn get_tool(&self, name: &str) -> Option { @@ -402,33 +514,25 @@ impl ServerHandler for FakeMcpServer { return Err(McpError::invalid_params(format!("unknown tool: {}", request.name), None)); }; - if !response.delay.is_zero() { - tokio::time::sleep(response.delay).await; + if let Some(gate) = &response.gate { + gate.pass().await; } if let Some(token) = context.meta.get_progress_token() { - for FakeProgress { progress, total, message } in response.progress { - let mut notification = ProgressNotificationParam::new(token.clone(), progress); - if let Some(total) = total { - notification = notification.with_total(total); + let notifications: Vec<_> = + response.progress.iter().map(|progress| progress.notification(&token)).collect(); + let peer = context.peer.clone(); + let notify = async move { + for notification in notifications { + let _ = peer.notify_progress(notification).await; } - if let Some(message) = message { - notification = notification.with_message(message); - } - let _ = context.peer.notify_progress(notification).await; - } - if !response.task_progress.is_empty() { - let peer = context.peer.clone(); - let token = token.clone(); + }; + if matches!(response.response, CallToolResponse::Task(_)) { tokio::spawn(async move { tokio::task::yield_now().await; - for (progress, total) in response.task_progress { - let mut notification = ProgressNotificationParam::new(token.clone(), progress); - if let Some(total) = total { - notification = notification.with_total(total); - } - let _ = peer.notify_progress(notification).await; - } + notify.await; }); + } else { + notify.await; } } Ok(response.response) @@ -440,6 +544,11 @@ type ToolHandler = Arc FakeToolResponse + Send #[derive(Default)] struct FakeMcpStateInner { tools: BTreeMap, + instructions: Option, + prompts: BTreeMap, + prompt_gate: Option, + tool_list_gate: Option, + prompt_elicitations: Option>, calls: Vec, client_capabilities: Option, tasks: HashMap>, @@ -448,8 +557,29 @@ struct FakeMcpStateInner { task_cancel_ids: Vec, task_get_failures: usize, task_update_failures: usize, + tool_list_requests: usize, tool_list_failures: usize, - peers: Vec>, + client: Option>, +} + +#[derive(Clone)] +struct FakeProgress { + progress: f64, + total: Option, + message: Option, +} + +impl FakeProgress { + fn notification(&self, token: &ProgressToken) -> ProgressNotificationParam { + let mut notification = ProgressNotificationParam::new(token.clone(), self.progress); + if let Some(total) = self.total { + notification = notification.with_total(total); + } + if let Some(message) = &self.message { + notification = notification.with_message(message.clone()); + } + notification + } } fn add_numbers() -> FakeTool { @@ -469,16 +599,6 @@ fn divide_numbers() -> FakeTool { }) } -fn slow_tool() -> FakeTool { - FakeTool::new("slow_tool") - .description("A tool that sleeps for a specified duration (for testing timeouts)") - .responds_with(|request| { - let sleep_ms = int_arg(request, "sleep_ms").unsigned_abs(); - FakeToolResponse::new(CallToolResult::structured(json!({ "message": format!("Slept for {sleep_ms}ms") }))) - .delay(Duration::from_millis(sleep_ms)) - }) -} - fn int_arg(request: &CallToolRequestParams, name: &str) -> i64 { request.arguments.as_ref().and_then(|args| args.get(name)).and_then(serde_json::Value::as_i64).unwrap_or_default() } diff --git a/crates/mcp-utils/src/testing/fake_task.rs b/crates/mcp-utils/src/testing/fake_task.rs new file mode 100644 index 000000000..67438cd36 --- /dev/null +++ b/crates/mcp-utils/src/testing/fake_task.rs @@ -0,0 +1,138 @@ +use super::{FakeMcpServer, FakeTool, FakeToolResponse, Gate, completed_task_payload}; +use rmcp::model::{ + CallToolResponse, CallToolResult, ContentBlock, CreateTaskResult, DetailedTask, ElicitRequest, ElicitRequestParams, + ElicitResult, ElicitationAction, InputRequest, InputRequests, InputRequiredResult, Task, TaskPayload, TaskStatus, +}; +use serde_json::json; + +#[derive(Clone)] +pub struct FakeTask { + seed: Task, + states: Vec, + progress: Vec<(f64, Option)>, + gate: Option, + get_failures: usize, + update_failures: usize, +} + +impl FakeTask { + pub fn new(task_id: &str) -> Self { + Self { + seed: task(task_id, TaskStatus::Working), + states: vec![working_task(task_id)], + progress: Vec::new(), + gate: None, + get_failures: 0, + update_failures: 0, + } + } + + pub fn seed(mut self, seed: Task) -> Self { + self.seed = seed; + self + } + + pub fn states(mut self, states: impl IntoIterator) -> Self { + self.states = states.into_iter().collect(); + self + } + + pub fn progress(mut self, progress: f64, total: Option) -> Self { + self.progress.push((progress, total)); + self + } + + pub fn gated(mut self, gate: Gate) -> Self { + self.gate = Some(gate); + self + } + + pub fn failing_gets(mut self, failures: usize) -> Self { + self.get_failures = failures; + self + } + + pub fn failing_updates(mut self, failures: usize) -> Self { + self.update_failures = failures; + self + } + + pub fn server(self, tool: &str) -> FakeMcpServer { + let task_id = self.seed.task_id.clone(); + let mut response = FakeToolResponse::task(CreateTaskResult::new(self.seed)); + for (progress, total) in self.progress { + response = response.progress(progress, total); + } + if let Some(gate) = self.gate { + response = response.gated(gate); + } + FakeMcpServer::new() + .with_tool(FakeTool::new(tool).responds(response)) + .with_task(task_id, self.states) + .with_task_get_failures(self.get_failures) + .with_task_update_failures(self.update_failures) + } +} + +pub fn task(task_id: &str, status: TaskStatus) -> Task { + let now = chrono::Utc::now().to_rfc3339(); + Task::new(task_id, status, now.clone(), now).with_poll_interval_ms(10) +} + +pub fn working_task(task_id: &str) -> DetailedTask { + DetailedTask::new(task(task_id, TaskStatus::Working), TaskPayload::Working) +} + +pub fn completed_task(task_id: &str, text: &str) -> DetailedTask { + let result = CallToolResult::success(vec![ContentBlock::text(text)]); + DetailedTask::new(task(task_id, TaskStatus::Completed), completed_task_payload(result)) +} + +pub fn failed_task(task_id: &str, message: &str) -> DetailedTask { + let error = serde_json::from_value(json!({ "code": "boom", "message": message })).expect("an error object"); + DetailedTask::new(task(task_id, TaskStatus::Failed), TaskPayload::Failed { error }) +} + +pub fn cancelled_task(task_id: &str) -> DetailedTask { + DetailedTask::new(task(task_id, TaskStatus::Cancelled), TaskPayload::Cancelled) +} + +pub fn input_required_task(task_id: &str, keys: &[&str]) -> DetailedTask { + let input_requests = keys.iter().map(|key| ((*key).to_string(), form_request("Provide a name"))).collect(); + DetailedTask::new(task(task_id, TaskStatus::InputRequired), TaskPayload::InputRequired { input_requests }) +} + +pub fn name_input(state: &str) -> CallToolResponse { + InputRequiredResult::new(Some(input_requests("name", form_request("Name?"))), Some(state.to_string())).into() +} + +pub fn accept_with_name(name: &str) -> ElicitResult { + ElicitResult::new(ElicitationAction::Accept).with_content(json!({ "name": name })) +} + +pub fn input_requests(key: &str, request: InputRequest) -> InputRequests { + InputRequests::from([(key.to_string(), request)]) +} + +pub fn form_request(message: &str) -> InputRequest { + let requested_schema = serde_json::from_value(json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"] + })) + .expect("a valid elicitation schema"); + InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::FormElicitationParams { + meta: None, + message: message.to_string(), + requested_schema, + })) +} + +pub fn url_request(message: &str, url: &str, elicitation_id: &str) -> InputRequest { + InputRequest::Elicitation(ElicitRequest::new(ElicitRequestParams::UrlElicitationParams { + meta: None, + message: message.to_string(), + url: url.to_string(), + elicitation_id: elicitation_id.to_string(), + })) +} diff --git a/crates/mcp-utils/src/testing/gate.rs b/crates/mcp-utils/src/testing/gate.rs new file mode 100644 index 000000000..13934397c --- /dev/null +++ b/crates/mcp-utils/src/testing/gate.rs @@ -0,0 +1,40 @@ +use std::sync::Arc; +use tokio::sync::watch; + +#[derive(Clone)] +pub struct Gate { + state: Arc>, +} + +#[derive(Clone, Copy, Default)] +struct GateState { + reached: bool, + open: bool, +} + +impl Gate { + pub fn new() -> Self { + Self { state: Arc::new(watch::Sender::new(GateState::default())) } + } + + pub fn open(&self) { + self.state.send_modify(|state| state.open = true); + } + + pub async fn reached(&self) { + let mut state = self.state.subscribe(); + let _ = state.wait_for(|state| state.reached).await; + } + + pub(crate) async fn pass(&self) { + self.state.send_modify(|state| state.reached = true); + let mut state = self.state.subscribe(); + let _ = state.wait_for(|state| state.open).await; + } +} + +impl Default for Gate { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/mcp-utils/src/tool_gateway/mod.rs b/crates/mcp-utils/src/tool_gateway/mod.rs deleted file mode 100644 index c6d193365..000000000 --- a/crates/mcp-utils/src/tool_gateway/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -mod transport; - -pub use transport::{UnixSocketMcpTransport, UnixSocketPath, UnixSocketServer, UnixSocketTransportError, connect}; - -/// Environment variable used to pass the active session endpoint to child tools. -pub const AETHER_MCP_IPC_SOCKET: &str = "AETHER_MCP_IPC_SOCKET"; - -/// Private adapter tool used by the progressive CLI to discover deferred servers. -pub const LIST_SERVERS_TOOL: &str = "_aether_list_servers"; diff --git a/crates/mcp-utils/src/tool_gateway/transport.rs b/crates/mcp-utils/src/tool_gateway/transport.rs deleted file mode 100644 index e94d9e0f6..000000000 --- a/crates/mcp-utils/src/tool_gateway/transport.rs +++ /dev/null @@ -1,294 +0,0 @@ -use rmcp::transport::async_rw::AsyncRwTransport; -use rmcp::{RoleClient, RoleServer, ServiceExt}; -use std::env::{temp_dir, var_os}; -use std::fs::{Permissions, create_dir_all, remove_dir}; -use std::fs::{remove_file, set_permissions}; -use std::{ - io, - os::unix::fs::PermissionsExt, - os::unix::net::UnixListener as StdUnixListener, - path::{Path, PathBuf}, -}; -use tokio::io::{ReadHalf, WriteHalf}; -use tokio::net::{UnixListener, UnixStream}; -use tokio::task::JoinHandle; -use tokio_util::sync::CancellationToken; -use uuid::Uuid; - -#[derive(Debug, thiserror::Error)] -pub enum UnixSocketTransportError { - #[error("failed to create MCP socket directory: {0}")] - CreateDirectory(#[source] io::Error), - #[error("failed to bind MCP socket: {0}")] - Bind(#[source] io::Error), - #[error("MCP socket path must be absolute")] - NotAbsolute, - #[error("MCP socket path is not valid UTF-8")] - InvalidPath, -} - -/// A session endpoint allocated by this process, or inherited from the session environment. -#[derive(Debug)] -pub struct UnixSocketPath { - directory: PathBuf, - socket: PathBuf, - /// Only the allocating process owns endpoint removal; inherited paths never remove. - remove_on_drop: bool, -} - -impl UnixSocketPath { - pub fn new() -> Result { - let short_id = &Uuid::new_v4().simple().to_string()[..8]; - let runtime_dir = - var_os("XDG_RUNTIME_DIR").filter(|path| Path::new(path).is_absolute()).map_or_else(temp_dir, PathBuf::from); - let socket_dir = runtime_dir.join("aether").join(format!("aether-{short_id}")); - let socket = socket_dir.join("ipc.sock"); - create_dir_all(&socket_dir).map_err(UnixSocketTransportError::CreateDirectory)?; - set_permissions(&socket_dir, Permissions::from_mode(0o700)) - .map_err(UnixSocketTransportError::CreateDirectory)?; - Ok(Self { directory: socket_dir, socket, remove_on_drop: true }) - } - - pub fn from_path(path: impl Into) -> Result { - let socket = path.into(); - if !socket.is_absolute() { - return Err(UnixSocketTransportError::NotAbsolute); - } - let directory = socket.parent().ok_or(UnixSocketTransportError::InvalidPath)?.to_path_buf(); - Ok(Self { directory, socket, remove_on_drop: false }) - } - - pub fn path(&self) -> &Path { - &self.socket - } - - pub fn directory(&self) -> &Path { - &self.directory - } -} - -impl Drop for UnixSocketPath { - fn drop(&mut self) { - if self.remove_on_drop { - let _ = remove_file(&self.socket); - let _ = remove_dir(&self.directory); - } - } -} - -pub struct UnixSocketMcpTransport { - path: UnixSocketPath, - listener: UnixListener, -} - -impl UnixSocketMcpTransport { - pub fn bind(path: UnixSocketPath) -> Result { - let _ = remove_file(path.path()); - let listener = StdUnixListener::bind(path.path()).map_err(UnixSocketTransportError::Bind)?; - listener.set_nonblocking(true).map_err(UnixSocketTransportError::Bind)?; - let listener = UnixListener::from_std(listener).map_err(UnixSocketTransportError::Bind)?; - Ok(Self { path, listener }) - } - - pub fn path(&self) -> &Path { - self.path.path() - } - - /// Serve connections until the returned server is dropped. - pub fn spawn(self, server: T) -> UnixSocketServer - where - T: Clone + ServiceExt + Send + 'static, - { - let Self { path, listener } = self; - let cancellation = CancellationToken::new(); - let accept_cancellation = cancellation.clone(); - let task = tokio::spawn(async move { - loop { - let accepted = tokio::select! { - () = accept_cancellation.cancelled() => break, - result = listener.accept() => result, - }; - let Ok((stream, _)) = accepted else { break }; - let server = server.clone(); - let connection_cancellation = accept_cancellation.clone(); - tokio::spawn(async move { - match server.serve(stream).await { - Ok(running) => { - let service_cancellation = running.cancellation_token(); - let mut waiting = Box::pin(running.waiting()); - tokio::select! { - () = connection_cancellation.cancelled() => service_cancellation.cancel(), - _ = &mut waiting => {} - } - } - Err(error) => tracing::debug!(%error, "MCP Unix socket client ended during initialization"), - } - }); - } - }); - UnixSocketServer { path, cancellation, task } - } -} - -/// Owns the accept task, connection cancellation, and the session endpoint. -/// Dropping it cancels in-flight connections and removes the endpoint. -pub struct UnixSocketServer { - path: UnixSocketPath, - cancellation: CancellationToken, - task: JoinHandle<()>, -} - -impl UnixSocketServer { - pub fn path(&self) -> &Path { - self.path.path() - } -} - -impl Drop for UnixSocketServer { - fn drop(&mut self) { - self.cancellation.cancel(); - self.task.abort(); - } -} - -/// Connect an rmcp client to an inherited session endpoint. -pub async fn connect( - path: impl AsRef, -) -> io::Result, WriteHalf>> { - let stream = UnixStream::connect(path).await?; - let (read, write) = tokio::io::split(stream); - Ok(AsyncRwTransport::new_client(read, write)) -} - -#[cfg(test)] -mod tests { - use super::*; - use rmcp::ServerHandler; - use rmcp::handler::server::router::tool::ToolRouter; - use rmcp::model::{ServerCapabilities, ServerConfig}; - use rmcp::{tool, tool_handler, tool_router}; - - #[derive(Clone)] - struct TestServer { - tool_router: ToolRouter, - } - - #[tool_router(allow_empty)] - impl TestServer {} - - #[allow(clippy::unused_async_trait_impl)] - #[tool_handler(router = self.tool_router)] - impl ServerHandler for TestServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - } - } - - #[test] - fn allocated_endpoint_is_removed_on_drop_without_binding() { - let path = UnixSocketPath::new().unwrap(); - let directory = path.directory().to_path_buf(); - assert!(directory.exists()); - drop(path); - assert!(!directory.exists()); - } - - #[test] - fn inherited_endpoint_is_not_removed_on_drop() { - let path = UnixSocketPath::new().unwrap(); - let directory = path.directory().to_path_buf(); - let inherited = UnixSocketPath::from_path(path.path()).unwrap(); - drop(inherited); - assert!(directory.exists()); - drop(path); - assert!(!directory.exists()); - } - - #[tokio::test] - async fn allocated_endpoint_is_private_and_removed_with_transport() { - let path = UnixSocketPath::new().unwrap(); - let directory = path.directory().to_path_buf(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - assert_eq!(std::fs::metadata(&directory).unwrap().permissions().mode() & 0o777, 0o700); - assert!(transport.path().exists()); - drop(transport); - assert!(!directory.exists()); - } - - #[tokio::test] - async fn spawned_endpoint_is_removed_with_server() { - let path = UnixSocketPath::new().unwrap(); - let directory = path.directory().to_path_buf(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - let socket = transport.path().to_path_buf(); - let server = transport.spawn(TestServer { tool_router: TestServer::tool_router() }); - assert!(socket.exists()); - drop(server); - assert!(!directory.exists()); - } - - #[tokio::test] - async fn connected_client_completes_initialization() { - let path = UnixSocketPath::new().unwrap(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - let socket = transport.path().to_path_buf(); - let _server = transport.spawn(TestServer { tool_router: TestServer::tool_router() }); - let _client = ().serve(connect(&socket).await.unwrap()).await.unwrap(); - } - - #[tokio::test] - async fn connected_client_can_list_empty_tools() { - let path = UnixSocketPath::new().unwrap(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - let socket = transport.path().to_path_buf(); - let _server = transport.spawn(TestServer { tool_router: TestServer::tool_router() }); - let client = ().serve(connect(&socket).await.unwrap()).await.unwrap(); - - assert_eq!(client.list_all_tools().await.unwrap().len(), 0); - } - - #[tokio::test] - async fn dropping_server_cancels_in_flight_connections() { - use rmcp::model::CallToolRequestParams; - use tokio::sync::watch; - - #[derive(Clone)] - struct SlowServer { - tool_router: ToolRouter, - started: watch::Sender, - } - - #[tool_router] - impl SlowServer { - #[tool(description = "Blocks until the connection is aborted")] - async fn slow(&self) -> String { - let _ = self.started.send(true); - std::future::pending::<()>().await; - "done".to_string() - } - } - - #[allow(clippy::unused_async_trait_impl)] - #[tool_handler(router = self.tool_router)] - impl ServerHandler for SlowServer { - fn get_info(&self) -> ServerConfig { - ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) - } - } - - let path = UnixSocketPath::new().unwrap(); - let transport = UnixSocketMcpTransport::bind(path).unwrap(); - let socket = transport.path().to_path_buf(); - let (started_tx, mut started_rx) = watch::channel(false); - let server = transport.spawn(SlowServer { tool_router: SlowServer::tool_router(), started: started_tx }); - - let client = ().serve(connect(&socket).await.unwrap()).await.unwrap(); - let call = tokio::spawn(async move { - let _ = client.call_tool_once(CallToolRequestParams::new("slow")).await; - }); - started_rx.changed().await.unwrap(); - - drop(server); - call.await.unwrap(); - } -} diff --git a/crates/mcp-utils/tests/client.rs b/crates/mcp-utils/tests/client.rs new file mode 100644 index 000000000..ec3ba7ecc --- /dev/null +++ b/crates/mcp-utils/tests/client.rs @@ -0,0 +1,788 @@ +use futures::StreamExt; +use futures::future::Either; +use mcp_utils::client::{ + CancellationToken, ClientError, ClientOptions, Elicitation, ElicitationSink, TaskErrorReason, ToolCallError, + ToolCallEvent, ToolCallOptions, +}; +use mcp_utils::model::{ + CallToolRequestParams, CallToolResponse, CallToolResult, CustomNotification, DEFAULT_MRTR_MAX_ROUNDS, DetailedTask, + DiscoverResult, ElicitRequestParams, ElicitResult, ElicitationAction, ElicitationCapability, ErrorCode, ErrorData, + FormElicitationCapability, Implementation, InitializeRequestParams, InitializeResult, InputRequests, + InputRequiredResult, ListToolsResult, PaginatedRequestParams, ProtocolVersion, ServerCapabilities, ServerConfig, + ServerNotification, Task, TaskPayload, TaskStatus, Tool, UrlElicitationCapability, +}; +use mcp_utils::testing::{ + ElicitationScript, FakeMcpServer, FakeMcpState, FakeTask, FakeTool, FakeToolResponse, accept_with_name, call, + cancelled_task, completed_task, connect, elicitation_host, failed_task, input_requests, input_required_task, + name_input, task, url_request, working_task, +}; +use rmcp::{ + RoleServer, ServerHandler, + service::{NotificationContext, RequestContext}, +}; +use serde_json::{Map, json}; +use std::borrow::Cow; +use std::future::{Ready, ready}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::{mpsc, watch}; +use tokio::time::Instant; + +#[tokio::test] +async fn connect_prefers_stateless_discovery_for_modern_servers() { + let negotiated = Negotiated::default(); + let client = connect("versioned", ModernOnlyServer(negotiated.clone()), &ClientOptions::default()).await; + + client.list_tools().await.expect("list tools"); + + assert_eq!(negotiated.observed().await, ProtocolVersion::V_2026_07_28); +} + +#[tokio::test] +async fn connect_selects_an_older_mutually_supported_revision() { + let negotiated = Negotiated::default(); + let client = connect("versioned", June2025Server(negotiated.clone()), &ClientOptions::default()).await; + + client.list_tools().await.expect("list tools"); + + assert_eq!(negotiated.observed().await, ProtocolVersion::V_2025_06_18); +} + +#[tokio::test] +async fn connect_falls_back_to_legacy_initialization() { + let negotiated = Negotiated::default(); + + let _client = connect("versioned", LegacyOnlyServer(negotiated.clone()), &ClientOptions::default()).await; + + assert_eq!(negotiated.observed().await, ProtocolVersion::V_2025_11_25); +} + +#[tokio::test] +async fn exposes_the_server_name_description_and_instructions() { + let server = FakeMcpServer::math().with_instructions("Use the math tools"); + let client = connect("fake", server, &ClientOptions::default()).await; + + assert_eq!(client.name(), "fake"); + assert_eq!(client.description().as_deref(), Some("A fake MCP server for testing")); + assert_eq!(client.instructions().as_deref(), Some("Use the math tools")); + assert!(client.list_tools().await.unwrap().iter().any(|tool| tool.name == "add_numbers")); +} + +#[tokio::test] +async fn servers_without_prompts_list_none() { + let client = connect("fake", FakeMcpServer::new(), &ClientOptions::default()).await; + + assert_eq!(client.list_prompts().await.unwrap().len(), 0); +} + +#[tokio::test] +async fn tool_calls_resolve_to_their_result() { + let client = connect("fake", FakeMcpServer::math(), &ClientOptions::default()).await; + + let result = call(&client, "add_numbers", json!({"a": 2, "b": 3})).await; + + assert_eq!(result.unwrap().structured_content, Some(json!({"sum": 5}))); +} + +#[tokio::test] +async fn tool_lists_start_with_the_current_tools_then_follow_each_change() { + let server = FakeMcpServer::math(); + let state = server.state(); + let client = connect("fake", server, &ClientOptions::default()).await; + let mut tool_lists = Box::pin(client.tool_lists()); + + let current = tool_names(tool_lists.next().await); + state.add_tool_and_notify(FakeTool::new("added_later")).await; + let changed = tool_names(tool_lists.next().await); + + assert_eq!(current, ["add_numbers", "divide_numbers"]); + assert_eq!(changed, ["add_numbers", "added_later", "divide_numbers"]); +} + +#[tokio::test] +async fn closing_any_clone_ends_the_session_for_all() { + let client = connect("fake", FakeMcpServer::new(), &ClientOptions::default()).await; + + client.clone().close().await; + + assert!(client.list_tools().await.is_err()); +} + +#[tokio::test] +async fn tasks_are_always_advertised_but_elicitation_needs_an_event_sink() { + let without_sink = advertised(ClientOptions::default()).await; + let (events, _host) = mpsc::channel(1); + let with_sink = + advertised(ClientOptions::default().elicitation(ElicitationSink::form_and_url(events.clone()))).await; + let url = ElicitationCapability::new().with_url(UrlElicitationCapability::new()); + let url_only = + advertised(ClientOptions::default().elicitation(ElicitationSink::new(events.clone(), url.clone()))).await; + let neither = + advertised(ClientOptions::default().elicitation(ElicitationSink::new(events, ElicitationCapability::new()))) + .await; + + assert!(without_sink.tasks && with_sink.tasks && url_only.tasks && neither.tasks); + assert_eq!(without_sink.elicitation, None); + assert_eq!(with_sink.elicitation, Some(url.with_form(FormElicitationCapability::new()))); + assert_eq!(url_only.elicitation, Some(ElicitationCapability::new().with_url(UrlElicitationCapability::new()))); + assert_eq!(neither.elicitation, None); +} + +#[tokio::test] +async fn elicitations_resolve_as_cancel_without_an_event_sink() { + let test = ElicitationTest::new(); + let client = connect("fake", test.server(), &ClientOptions::default()).await; + + call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(test.answer(), json!({"action": "cancel"})); +} + +#[tokio::test] +async fn dropping_an_elicitation_answers_cancel() { + let test = ElicitationTest::new(); + let (options, mut host) = elicitation_host(); + let client = connect("fake", test.server(), &options).await; + let host = tokio::spawn(async move { + let Some(Elicitation::Request(elicitation)) = host.recv().await else { panic!("expected an elicitation") }; + drop(elicitation); + }); + + call(&client, "ask", json!({})).await.unwrap(); + host.await.unwrap(); + + assert_eq!(test.answer(), json!({"action": "cancel"})); +} + +#[tokio::test] +async fn elicitation_responses_reach_the_server_and_name_their_source() { + let test = ElicitationTest::new(); + let (options, mut host) = elicitation_host(); + let client = connect("fake", test.server(), &options).await; + let host = tokio::spawn(async move { + let Some(Elicitation::Request(elicitation)) = host.recv().await else { panic!("expected an elicitation") }; + let server = elicitation.server.clone(); + elicitation.respond(ElicitResult::new(ElicitationAction::Accept).with_content(json!({"name": "Ferris"}))); + server + }); + + call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(host.await.unwrap(), "fake"); + assert_eq!(test.answer(), json!({"action": "accept", "content": {"name": "Ferris"}})); +} + +#[tokio::test] +async fn elicitations_answer_cancel_once_the_host_stops_listening() { + let test = ElicitationTest::new(); + let (options, host) = elicitation_host(); + let client = connect("fake", test.server(), &options).await; + drop(host); + + call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(test.answer(), json!({"action": "cancel"})); +} + +#[tokio::test] +async fn elicitation_completion_notifications_name_their_source() { + let (options, mut host) = elicitation_host(); + let client = connect("linear", CompletionServer, &options).await; + + call(&client, "complete", json!({})).await.unwrap(); + + let event = host.recv().await.expect("completion event"); + assert!(matches!(event, Elicitation::Complete { server, id } if server == "linear" && id == "oauth-1")); +} + +#[tokio::test] +async fn input_required_without_requests_or_state_aborts() { + let server = asking(CallToolResponse::from(InputRequiredResult::from_input_requests(InputRequests::new()))); + let client = connect("fake", server, &ClientOptions::default()).await; + + let result = call(&client, "ask", json!({})).await; + + assert!(matches!(result, Err(ToolCallError::EmptyInputRequired)), "{result:?}"); +} + +#[tokio::test] +async fn state_only_rounds_poll_until_the_server_completes() { + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask") + .responds(poll("first")) + .when_state("first", poll("second")) + .when_state("second", FakeToolResponse::text("done")), + ); + let state = server.state(); + let client = connect("fake", server, &ClientOptions::default()).await; + + let result = call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(text(&result), "done"); + assert_eq!(request_states(&state), [None, Some("first".to_string()), Some("second".to_string())]); +} + +#[tokio::test(start_paused = true)] +async fn state_only_rounds_back_off_exponentially_up_to_a_cap_and_reset_after_input() { + let mut tool = FakeTool::new("ask").responds(poll("1")); + for round in 1..7 { + tool = tool.when_state(round.to_string(), poll(&(round + 1).to_string())); + } + let tool = tool + .when_state("7", name_input("asked")) + .when_state("asked", poll("after")) + .when_state("after", FakeToolResponse::text("done")); + let server = FakeMcpServer::new().with_tool(tool); + let state = server.state(); + let (options, _user) = ElicitationScript::answering([accept_with_name("Ferris")]); + let client = connect("fake", server, &options).await; + + call(&client, "ask", json!({})).await.unwrap(); + + let arrivals = state.calls_for("ask").into_iter().map(|call| call.at).collect::>(); + let gaps = arrivals.windows(2).map(|pair| (pair[1] - pair[0]).as_millis()).collect::>(); + assert_eq!(gaps, [50, 100, 200, 400, 800, 1600, 1600, 0, 50]); +} + +#[tokio::test] +async fn the_server_may_not_ask_again_once_the_user_cancels() { + let server = FakeMcpServer::new() + .with_tool(FakeTool::new("ask").responds(name_input("asked")).when_state("asked", name_input("asked"))); + let client = connect("fake", server, &ClientOptions::default()).await; + + let result = call(&client, "ask", json!({})).await; + + assert!(matches!(result, Err(ToolCallError::RePromptAfterCancel)), "{result:?}"); +} + +#[tokio::test] +async fn the_server_may_still_finish_by_polling_once_the_user_cancels() { + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask") + .responds(name_input("asked")) + .when_state("asked", poll("finishing")) + .when_state("finishing", FakeToolResponse::text("done")), + ); + let client = connect("fake", server, &ClientOptions::default()).await; + + let result = call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(text(&result), "done"); +} + +#[tokio::test] +async fn answered_input_rounds_continue_up_to_the_round_limit() { + let server = FakeMcpServer::new() + .with_tool(FakeTool::new("ask").responds(name_input("asked")).when_state("asked", name_input("asked"))); + let state = server.state(); + let (options, _user) = + ElicitationScript::answering(vec![ElicitResult::new(ElicitationAction::Accept); DEFAULT_MRTR_MAX_ROUNDS]); + let client = connect("fake", server, &options).await; + + let result = call(&client, "ask", json!({})).await; + + assert!(matches!(result, Err(ToolCallError::InputRoundsExceeded)), "{result:?}"); + assert_eq!(state.calls_for("ask").len(), DEFAULT_MRTR_MAX_ROUNDS + 1); +} + +#[tokio::test] +async fn form_answers_and_request_state_are_echoed_back_verbatim() { + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask") + .responds(name_input(TRICKY_STATE)) + .when_state(TRICKY_STATE, FakeToolResponse::text("done")), + ); + let state = server.state(); + let (options, _user) = ElicitationScript::answering([accept_with_name("Ferris")]); + let client = connect("fake", server, &options).await; + + let result = call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(text(&result), "done"); + let calls = state.calls_for("ask"); + assert_eq!(calls[1].request.request_state.as_deref(), Some(TRICKY_STATE)); + assert_eq!( + calls[1].request.input_responses.as_ref().unwrap()["name"], + json!({"action": "accept", "content": {"name": "Ferris"}}) + ); +} + +#[tokio::test] +async fn url_elicitations_are_answered_once_then_polled_to_completion() { + let consent = input_requests("consent", url_request("Authorize", "https://example.com/auth", "el-1")); + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask") + .responds(InputRequiredResult::new(Some(consent), Some("url-opened".to_string()))) + .when_state("url-opened", poll("polling")) + .when_state("polling", FakeToolResponse::text("done")), + ); + let state = server.state(); + let (options, user) = ElicitationScript::answering([ElicitResult::new(ElicitationAction::Accept)]); + let client = connect("fake", server, &options).await; + + let result = call(&client, "ask", json!({})).await.unwrap(); + + assert_eq!(text(&result), "done"); + assert_eq!(request_states(&state), [None, Some("url-opened".to_string()), Some("polling".to_string())]); + let asked = user.captured(); + assert_eq!(asked.len(), 1, "polling must not prompt again"); + assert!(matches!( + &asked[0].request, + ElicitRequestParams::UrlElicitationParams { url, elicitation_id, .. } + if url == "https://example.com/auth" && elicitation_id == "el-1" + )); +} + +#[tokio::test] +async fn progress_is_reported_across_input_rounds() { + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask") + .responds(FakeToolResponse::new(name_input("asked")).progress(1.0, Some(2.0))) + .when_state("asked", FakeToolResponse::text("done").progress(2.0, Some(2.0))), + ); + let (options, _user) = ElicitationScript::answering([accept_with_name("Ferris")]); + let client = connect("fake", server, &options).await; + + let events = client.call_tool("ask", Map::new(), ToolCallOptions::default()).collect::>().await; + + assert_eq!(progress(&events), [1.0, 2.0]); + assert!(matches!(events.last(), Some(ToolCallEvent::Done(Ok(result))) if text(result) == "done")); +} + +#[tokio::test] +async fn created_tasks_are_polled_until_they_complete() { + let result = task_test([working_task(TASK), completed_task(TASK, "finished")]).run().await; + + assert!(matches!(result.events.first(), Some(ToolCallEvent::TaskCreated(created)) if created.task.task_id == TASK)); + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Ok(result))) if text(result) == "finished")); + assert_eq!(result.state.task_get_ids(), [TASK, TASK]); +} + +#[tokio::test] +async fn intermediate_task_statuses_are_reported_before_the_result() { + let halfway = + DetailedTask::new(task(TASK, TaskStatus::Working).with_status_message("halfway"), TaskPayload::Working); + + let result = task_test([halfway, completed_task(TASK, "finished")]).run().await; + + let status = result.events.iter().position( + |event| matches!(event, ToolCallEvent::TaskStatus(task) if task.status_message.as_deref() == Some("halfway")), + ); + assert_eq!(status, Some(result.events.len() - 2)); +} + +#[tokio::test] +async fn a_terminal_seed_is_still_fetched_for_its_result() { + let result = + task_test([completed_task(TASK, "already done")]).with_task(task(TASK, TaskStatus::Completed)).run().await; + + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Ok(result))) if text(result) == "already done")); + assert_eq!(result.state.task_get_ids(), [TASK]); +} + +#[tokio::test] +async fn call_tool_forwards_progress_after_task_creation() { + let server = FakeTask::new(TASK) + .progress(1.0, Some(2.0)) + .states([working_task(TASK), completed_task(TASK, "finished")]) + .server("deferred"); + let client = connect("fake", server, &ClientOptions::default()).await; + + let events = client.call_tool("deferred", Map::new(), ToolCallOptions::default()).collect::>().await; + + assert!(matches!(events.first(), Some(ToolCallEvent::TaskCreated(_)))); + assert_eq!(progress(&events), [1.0]); + assert!(matches!(events.last(), Some(ToolCallEvent::Done(Ok(_))))); +} + +#[tokio::test] +async fn call_tool_handles_huge_task_ttl() { + let result = task_test([completed_task(TASK, "finished")]) + .with_task(task(TASK, TaskStatus::Working).with_ttl_ms(u64::MAX)) + .run() + .await; + + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Ok(_))))); +} + +#[tokio::test] +async fn call_tool_handles_huge_execution_timeout() { + let result = task_test([completed_task(TASK, "finished")]).with_timeout(Duration::MAX).run().await; + + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Ok(_))))); +} + +#[tokio::test] +async fn task_input_is_answered_through_tasks_update() { + let answer = accept_with_name("Ada"); + + let result = task_test([input_required_task(TASK, &["name"]), completed_task(TASK, "accepted")]) + .answering([answer.clone()]) + .run() + .await; + + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Ok(result))) if text(result) == "accepted")); + let updates = result.state.task_updates(); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].task_id, TASK); + assert_eq!(updates[0].input_responses.get("name"), Some(&serde_json::to_value(&answer).unwrap())); +} + +#[tokio::test] +async fn terminal_task_failures_name_their_reason() { + let malformed = DetailedTask::new( + task(TASK, TaskStatus::Completed), + TaskPayload::Completed { result: serde_json::from_value(json!({ "unexpected": true })).unwrap() }, + ); + + let failed = task_test([failed_task(TASK, "task exploded")]).run().await; + let cancelled = task_test([cancelled_task(TASK)]).run().await; + let malformed = task_test([malformed]).run().await; + + assert!(matches!(failed.failure(), TaskErrorReason::Failed { error } if error["message"] == "task exploded")); + assert!(matches!(cancelled.failure(), TaskErrorReason::Cancelled)); + assert!(matches!(malformed.failure(), TaskErrorReason::MalformedResult(_))); +} + +#[tokio::test] +async fn expired_tasks_fail_without_polling() { + let mut seed = task(TASK, TaskStatus::Working); + seed.created_at = "2020-01-01T00:00:00Z".to_string(); + seed.ttl_ms = Some(1); + + let result = task_test([completed_task(TASK, "too late")]).with_task(seed).run().await; + + assert!(matches!(result.failure(), TaskErrorReason::Expired)); + assert_eq!(result.state.task_get_ids(), Vec::::new()); +} + +#[tokio::test] +async fn a_failed_poll_abandons_the_task() { + let result = task_test([completed_task(TASK, "must not finish")]).failing_gets(1).run().await; + + assert!(matches!(result.failure(), TaskErrorReason::Get(_))); + assert_eq!(result.state.task_get_ids(), [TASK]); + assert_eq!(result.state.task_cancel_ids(), [TASK]); +} + +#[tokio::test] +async fn a_failed_input_update_abandons_the_task() { + let result = task_test([input_required_task(TASK, &["name"])]) + .answering([ElicitResult::new(ElicitationAction::Accept)]) + .failing_updates(1) + .run() + .await; + + assert!(matches!(result.failure(), TaskErrorReason::Update(_))); + assert_eq!(result.state.task_updates().len(), 1); + assert_eq!(result.state.task_cancel_ids(), [TASK]); +} + +#[tokio::test] +async fn call_tool_cancellation_cancels_server_task_and_ends_stream() { + let server = FakeTask::new(TASK).server("deferred"); + let state = server.state(); + let client = connect("fake", server, &ClientOptions::default()).await; + let cancel = CancellationToken::new(); + let options = ToolCallOptions { cancel: cancel.clone(), ..ToolCallOptions::default() }; + let mut events = client.call_tool("deferred", Map::new(), options); + + assert!(matches!(events.next().await, Some(ToolCallEvent::TaskCreated(_)))); + cancel.cancel(); + let mut last = None; + while let Some(event) = events.next().await { + last = Some(event); + } + + assert!(matches!(last, Some(ToolCallEvent::Done(Err(ToolCallError::Cancelled))))); + assert_eq!(state.task_cancel_ids(), [TASK]); +} + +#[tokio::test(start_paused = true)] +async fn call_tool_deadline_includes_task_elicitation() { + let result = task_test([input_required_task(TASK, &["name"])]).with_timeout(Duration::from_millis(25)).run().await; + + assert!(matches!(result.events.last(), Some(ToolCallEvent::Done(Err(ToolCallError::TimedOut(_)))))); + assert_eq!(result.state.task_get_ids(), [TASK]); + assert_eq!(result.state.task_cancel_ids(), [TASK]); +} + +#[tokio::test] +async fn a_round_repeating_any_answered_key_is_abandoned_without_a_partial_update() { + let result = task_test([input_required_task(TASK, &["name"]), input_required_task(TASK, &["name", "nickname"])]) + .answering([ElicitResult::new(ElicitationAction::Accept)]) + .run() + .await; + + assert!(matches!(result.failure(), TaskErrorReason::RepeatedInput)); + assert_eq!(result.state.task_updates().len(), 1); + assert_eq!(result.state.task_cancel_ids(), [TASK]); +} + +const TASK: &str = "task-1"; +const TRICKY_STATE: &str = "st.ate/with+special=chars and spaces \"quotes\"\n\ttab"; + +fn asking(response: impl Into) -> FakeMcpServer { + FakeMcpServer::new().with_tool(FakeTool::new("ask").responds(response)) +} + +fn poll(state: &str) -> CallToolResponse { + InputRequiredResult::new(None, Some(state.to_string())).into() +} + +fn progress(events: &[ToolCallEvent]) -> Vec { + events + .iter() + .filter_map(|event| match event { + ToolCallEvent::Progress(progress) => Some(progress.progress), + _ => None, + }) + .collect() +} + +fn request_states(state: &FakeMcpState) -> Vec> { + state.calls_for("ask").into_iter().map(|call| call.request.request_state).collect() +} + +fn text(result: &CallToolResult) -> &str { + result.content.first().and_then(|content| content.as_text()).map_or("", |text| text.text.as_str()) +} + +fn tool_names(tools: Option, ClientError>>) -> Vec { + let tools = tools.expect("a tool list").expect("the server lists its tools"); + let mut names = tools.into_iter().map(|tool| tool.name.to_string()).collect::>(); + names.sort(); + names +} + +struct TaskTest { + task: FakeTask, + timeout: Option, + answers: Option>, +} + +struct TaskTestResult { + events: Vec, + state: FakeMcpState, +} + +fn task_test(states: impl IntoIterator) -> TaskTest { + TaskTest { task: FakeTask::new(TASK).states(states), timeout: None, answers: None } +} + +impl TaskTest { + fn with_task(mut self, seed: Task) -> Self { + self.task = self.task.seed(seed); + self + } + + fn with_timeout(mut self, timeout: Duration) -> Self { + self.timeout = Some(timeout); + self + } + + fn answering(mut self, answers: impl IntoIterator) -> Self { + self.answers = Some(answers.into_iter().collect()); + self + } + + fn failing_gets(mut self, failures: usize) -> Self { + self.task = self.task.failing_gets(failures); + self + } + + fn failing_updates(mut self, failures: usize) -> Self { + self.task = self.task.failing_updates(failures); + self + } + + async fn run(self) -> TaskTestResult { + let server = self.task.server("deferred"); + let state = server.state(); + let (options, _host) = if let Some(answers) = self.answers { + let (options, script) = ElicitationScript::answering(answers); + (options, Either::Left(script)) + } else { + let (options, host) = elicitation_host(); + (options, Either::Right(host)) + }; + let client = connect("fake", server, &options).await; + + let options = ToolCallOptions { timeout: self.timeout, ..ToolCallOptions::default() }; + let events = client.call_tool("deferred", Map::new(), options).collect().await; + TaskTestResult { events, state } + } +} + +impl TaskTestResult { + fn failure(&self) -> &TaskErrorReason { + match self.events.last() { + Some(ToolCallEvent::Done(Err(ToolCallError::Task { task_id, reason }))) if task_id == TASK => reason, + last => panic!("expected the task to fail, got {last:?}"), + } + } +} + +struct Advertised { + tasks: bool, + elicitation: Option, +} + +async fn advertised(options: ClientOptions) -> Advertised { + let server = FakeMcpServer::new(); + let state = server.state(); + let _client = connect("fake", server, &options).await; + let capabilities = state.client_capabilities().expect("capabilities are sent during discovery"); + Advertised { + tasks: capabilities + .extensions + .as_ref() + .is_some_and(|extensions| extensions.contains_key("io.modelcontextprotocol/tasks")), + elicitation: capabilities.elicitation, + } +} + +struct ElicitationTest { + server: FakeMcpServer, +} + +impl ElicitationTest { + fn new() -> Self { + let server = FakeMcpServer::new().with_tool( + FakeTool::new("ask").responds(name_input("asked")).when_state("asked", FakeToolResponse::text("done")), + ); + Self { server } + } + + fn server(&self) -> FakeMcpServer { + self.server.clone() + } + + fn answer(&self) -> serde_json::Value { + let calls = self.server.state().calls_for("ask"); + calls.last().and_then(|call| call.request.input_responses.as_ref()).expect("the tool was answered")["name"] + .clone() + } +} + +#[derive(Clone)] +struct CompletionServer; + +impl ServerHandler for CompletionServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn call_tool( + &self, + _request: CallToolRequestParams, + context: RequestContext, + ) -> Result { + let completion = + CustomNotification::new("notifications/elicitation/complete", Some(json!({ "elicitationId": "oauth-1" }))); + context.peer.send_notification(ServerNotification::CustomNotification(completion)).await.ok(); + Ok(CallToolResult::success(Vec::new()).into()) + } +} + +#[derive(Clone)] +struct Negotiated(Arc>>); + +impl Default for Negotiated { + fn default() -> Self { + Self(Arc::new(watch::channel(None).0)) + } +} + +impl Negotiated { + async fn observed(&self) -> ProtocolVersion { + let mut observed = self.0.subscribe(); + observed.wait_for(Option::is_some).await.expect("recorder alive").clone().expect("checked by wait_for") + } + + fn record(&self, version: Option) { + self.0.send_replace(version); + } + + fn list_tools(&self, context: &RequestContext) -> Ready> { + self.record(context.protocol_version()); + ready(Ok(ListToolsResult::with_all_items(Vec::new()))) + } +} + +#[derive(Clone, Default)] +struct ModernOnlyServer(Negotiated); + +impl ServerHandler for ModernOnlyServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("modern-only", "1.0.0")) + .with_protocol_version(ProtocolVersion::V_2026_07_28) + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Owned(vec![ProtocolVersion::V_2026_07_28]) + } + + fn initialize( + &self, + _request: InitializeRequestParams, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + ready(Err(ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "initialize is not supported", None))) + } + + fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> impl Future> + Send + '_ { + self.0.list_tools(&context) + } +} + +#[derive(Clone, Default)] +struct June2025Server(Negotiated); + +impl ServerHandler for June2025Server { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("older-revision", "1.0.0")) + .with_protocol_version(ProtocolVersion::V_2025_06_18) + } + + fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> { + Cow::Owned(vec![ProtocolVersion::V_2025_06_18]) + } + + fn list_tools( + &self, + _request: Option, + context: RequestContext, + ) -> impl Future> + Send + '_ { + self.0.list_tools(&context) + } +} + +#[derive(Clone, Default)] +struct LegacyOnlyServer(Negotiated); + +impl ServerHandler for LegacyOnlyServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("legacy", "1.0.0")) + .with_protocol_version(ProtocolVersion::V_2025_11_25) + } + + fn discover( + &self, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + ready(Err(ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "server/discover is not supported", None))) + } + + fn on_initialized(&self, context: NotificationContext) -> impl Future + Send + '_ { + self.0.record(context.peer.peer_info().map(|info| info.protocol_version.clone())); + ready(()) + } +} diff --git a/crates/mcp-utils/tests/doc_examples.rs b/crates/mcp-utils/tests/doc_examples.rs deleted file mode 100644 index cbd4647f5..000000000 --- a/crates/mcp-utils/tests/doc_examples.rs +++ /dev/null @@ -1,10 +0,0 @@ -//! Ensures the JSON examples embedded in config doc comments stay valid. -//! These `.md` files are rendered into the auto-generated settings reference. - -use aether_doctest::assert_json_examples; -use mcp_utils::client::McpServerConfig; - -#[test] -fn mcp_server_config_examples_parse() { - assert_json_examples::("mcp_server_config.md", include_str!("../src/docs/mcp_server_config.md")); -} diff --git a/crates/mcp-utils/tests/elicitation_notifications.rs b/crates/mcp-utils/tests/elicitation_notifications.rs deleted file mode 100644 index f745c3545..000000000 --- a/crates/mcp-utils/tests/elicitation_notifications.rs +++ /dev/null @@ -1,38 +0,0 @@ -use mcp_utils::client::{McpClient, McpClientEvent, client_capabilities}; -use mcp_utils::testing::connect; -use rmcp::ServerHandler; -use rmcp::model::{ClientConfig, CustomNotification, Implementation, ServerNotification}; -use tokio::sync::mpsc; - -#[derive(Clone)] -struct CompletionServer; - -impl ServerHandler for CompletionServer {} - -#[tokio::test] -async fn custom_completion_notification_includes_the_source_server() { - let (event_tx, mut event_rx) = mpsc::channel(1); - let client = McpClient::new( - ClientConfig::new(client_capabilities(), Implementation::new("test-client", "1.0.0")), - "linear".to_string(), - event_tx, - ); - let (server, client) = connect(CompletionServer, client).await.expect("connect MCP peers"); - - server - .send_notification(ServerNotification::CustomNotification(CustomNotification::new( - "notifications/elicitation/complete", - Some(serde_json::json!({ "elicitationId": "oauth-1" })), - ))) - .await - .expect("send completion notification"); - - let event = event_rx.recv().await.expect("receive completion event"); - assert!(matches!( - event, - McpClientEvent::ElicitationComplete { server_name, elicitation_id } - if server_name == "linear" && elicitation_id == "oauth-1" - )); - - client.cancel().await.expect("close MCP client"); -} diff --git a/crates/mcp-utils/tests/gateway.rs b/crates/mcp-utils/tests/gateway.rs new file mode 100644 index 000000000..01a179c1a --- /dev/null +++ b/crates/mcp-utils/tests/gateway.rs @@ -0,0 +1,327 @@ +use futures::{FutureExt, StreamExt}; +use mcp_utils::client::{ClientError, ClientOptions, HeaderMap, OAuthSettings, ToolCall, ToolCallOptions, Transport}; +use mcp_utils::gateway::{ConnectedServer, GatewayError, McpGateway, ServerSpec, ServerState}; +use mcp_utils::model::{ContentBlock, GetPromptResult, Tool, ToolAnnotations}; +use mcp_utils::server::{McpServer, UnixSocket}; +use mcp_utils::testing::{FakeMcpServer, FakeMcpState, FakeTool, FakeToolResponse, args, elicitation_host, wait_until}; +use serde_json::{Map, Value, json}; +use std::collections::HashMap; +use std::sync::Weak; + +#[tokio::test] +async fn servers_are_connecting_until_their_handshake_completes() { + let gateway = gateway(vec![stdio("silent", "sleep", &["60"])]); + + let statuses = gateway.statuses(); + assert_eq!(statuses.len(), 1); + assert_eq!(statuses[0].state, ServerState::Connecting); + gateway.shutdown().await; +} + +#[tokio::test] +async fn ready_waits_until_every_server_connects_or_fails() { + let gateway = gateway(vec![ + ServerSpec::in_process("math", FakeMcpServer::math()), + stdio("broken", "aether-missing-mcp-server", &[]), + ]); + + gateway.ready().await; + let statuses = gateway.statuses(); + + assert_eq!(statuses.iter().map(|entry| entry.name.as_str()).collect::>(), ["math", "broken"]); + assert_eq!(statuses[0].state, ServerState::Connected { tool_count: 2 }); + assert!(matches!(&statuses[1].state, ServerState::Failed { error } if error.contains("Failed to spawn"))); + assert!(!statuses[1].can_sign_in); +} + +#[tokio::test] +async fn ready_resolves_immediately_without_servers() { + let gateway = ready(gateway(Vec::new())).await; + + assert_eq!(gateway.statuses().len(), 0); +} + +#[tokio::test] +async fn unreachable_http_servers_fail_without_offering_oauth() { + let (options, _host) = elicitation_host(); + let gateway = ready(McpGateway::new(vec![unreachable_http("remote")], &options)).await; + + let status = &gateway.statuses()[0]; + assert!(matches!(status.state, ServerState::Failed { .. })); + assert!(!status.can_sign_in); +} + +#[tokio::test] +async fn failing_servers_do_not_block_the_others() { + let servers = vec![ + unreachable_http("failing"), + ServerSpec::in_process("working", FakeMcpServer::new()), + unreachable_http("also_failing"), + ]; + let gateway = ready(gateway(servers)).await; + + let statuses = gateway.statuses(); + assert_eq!(gateway.servers().iter().map(ConnectedServer::name).collect::>(), ["working"]); + assert!(matches!(statuses[0].state, ServerState::Failed { .. })); + assert!(matches!(statuses[1].state, ServerState::Connected { .. })); + assert!(matches!(statuses[2].state, ServerState::Failed { .. })); +} + +#[tokio::test] +async fn servers_report_their_tools_as_the_server_named_them() { + let server = FakeMcpServer::math() + .with_tool(FakeTool::new("lookup").description("Looks up").annotations(ToolAnnotations::new().read_only(true))); + let gateway = ready(gateway(vec![ServerSpec::in_process("math", server)])).await; + + let servers = gateway.servers(); + + let tools = servers[0].tools(); + assert_eq!(names(tools), ["add_numbers", "divide_numbers", "lookup"]); + let lookup = &tools[2]; + assert_eq!(lookup.description.as_deref(), Some("Looks up")); + assert_eq!(Value::Object((*lookup.input_schema).clone()), json!({"type": "object", "properties": {}})); + assert_eq!(lookup.annotations.as_ref().and_then(|annotations| annotations.read_only_hint), Some(true)); +} + +#[tokio::test] +async fn calls_reach_the_server_by_the_name_it_gave_the_tool() { + let gateway = ready(gateway(vec![ServerSpec::in_process("math", FakeMcpServer::math())])).await; + + let result = + gateway.servers()[0].call_tool("add_numbers", args(json!({"a": 2, "b": 3})), ToolCallOptions::default()); + + assert_eq!(result.result().await.unwrap().structured_content, Some(json!({"sum": 5}))); +} + +#[tokio::test] +async fn tool_list_changes_refresh_the_tools() { + let server = FakeMcpServer::math(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("math", server)])).await; + + state.add_tool_and_notify(FakeTool::new("added_later")).await; + wait_until(&gateway, |gateway| names(gateway.servers()[0].tools()).contains(&"added_later".to_string())).await; + state.clear_tools_and_notify().await; + wait_until(&gateway, |gateway| gateway.servers()[0].tools().is_empty()).await; + + assert_eq!(gateway.statuses()[0].state, ServerState::Connected { tool_count: 0 }); +} + +#[tokio::test] +async fn failed_tool_refreshes_keep_the_previous_tools() { + let server = FakeMcpServer::math(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("math", server)])).await; + + fail_a_refresh(&state).await; + fail_a_refresh(&state).await; + + assert_eq!(names(gateway.servers()[0].tools()), ["add_numbers", "divide_numbers"]); + assert_eq!(gateway.statuses()[0].state, ServerState::Connected { tool_count: 2 }); +} + +#[tokio::test] +async fn changes_report_only_what_happens_after_subscribing() { + let server = FakeMcpServer::math(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("math", server)])).await; + let mut changes = gateway.changes(); + + assert_eq!(changes.next().now_or_never(), None); + state.add_tool_and_notify(FakeTool::new("added_later")).await; + assert_eq!(changes.next().await, Some(())); +} + +#[tokio::test] +async fn prompts_are_listed_with_their_server_and_expanded_by_it() { + let docs = FakeMcpServer::new().with_prompt("greet", "Hello {name}"); + let gateway = ready(gateway(vec![ + ServerSpec::in_process("docs", docs), + ServerSpec::in_process("math", FakeMcpServer::math()), + ])) + .await; + + let prompts = gateway.list_prompts().await; + let expanded = gateway.get_prompt("docs", "greet", args(json!({"name": "Ada"}))).await.unwrap(); + + let listed = prompts.iter().map(|listed| (listed.server.as_str(), listed.prompt.name.as_str())).collect::>(); + assert_eq!(listed, [("docs", "greet")]); + assert_eq!(prompt_text(&expanded), "Hello Ada"); + assert!(matches!( + gateway.get_prompt("missing", "greet", Map::new()).await, + Err(GatewayError::ServerNotFound(name)) if name == "missing" + )); +} + +#[tokio::test] +async fn authenticate_needs_a_known_oauth_server() { + let gateway = ready(gateway(vec![ServerSpec::in_process("math", FakeMcpServer::math())])).await; + + assert!( + matches!(gateway.authenticate("missing").await, Err(GatewayError::ServerNotFound(name)) if name == "missing") + ); + assert!( + matches!(gateway.authenticate("math").await, Err(GatewayError::Client(ClientError::OAuthUnavailable { server })) if server == "math") + ); + assert!(!gateway.statuses()[0].can_sign_in); +} + +#[tokio::test] +async fn shutdown_releases_in_process_servers_and_closes_the_gateway() { + let server = FakeMcpServer::new(); + let alive = server.alive(); + let gateway = ready(gateway(vec![ServerSpec::in_process("probe", server)])).await; + + gateway.shutdown().await; + + assert!(alive.upgrade().is_none()); + assert_eq!(gateway.statuses().len(), 0); + gateway.ready().await; + assert_eq!(gateway.statuses().len(), 0); + assert!(matches!(gateway.authenticate("probe").await, Err(GatewayError::Closed))); + assert!(gateway.servers().is_empty()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn shutdown_releases_in_process_servers_while_another_handle_is_held() { + let server = FakeMcpServer::new(); + let alive = server.alive(); + let gateway = ready(gateway(vec![ServerSpec::in_process("probe", server)])).await; + let _held = gateway.clone(); + + gateway.shutdown().await; + + released(&alive).await; +} + +#[tokio::test] +async fn dropping_every_handle_releases_in_process_servers() { + let server = FakeMcpServer::new(); + let alive = server.alive(); + let gateway = ready(gateway(vec![ServerSpec::in_process("probe", server)])).await; + let mut changes = gateway.changes(); + + drop(gateway); + while changes.next().await.is_some() {} + + released(&alive).await; +} + +#[tokio::test] +async fn shutdown_cancels_calls_in_flight() { + let server = slow_server(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("slow", server)])).await; + let in_flight = tokio::spawn(slow_call(&gateway.servers()[0]).result()); + started(&state).await; + + gateway.shutdown().await; + + assert!(in_flight.await.unwrap().is_err()); +} + +#[tokio::test] +async fn dropping_every_handle_stops_the_gateway_and_cancels_calls_in_flight() { + let server = slow_server(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("slow", server)])).await; + let in_flight = tokio::spawn(slow_call(&gateway.servers()[0]).result()); + let mut changes = gateway.changes(); + started(&state).await; + + drop(gateway); + while changes.next().await.is_some() {} + + assert!(in_flight.await.unwrap().is_err()); +} + +#[tokio::test] +async fn refreshing_an_unchanged_tool_list_reports_no_change() { + let server = FakeMcpServer::math(); + let state = server.state(); + let gateway = ready(gateway(vec![ServerSpec::in_process("math", server)])).await; + let mut changes = gateway.changes(); + + refresh(&state).await; + state.fail_next_tool_list(); + refresh(&state).await; + + assert!(changes.next().now_or_never().is_none(), "an unchanged tool list was reported as a change"); +} + +#[tokio::test] +async fn a_server_that_disconnects_is_reported_as_failed() { + let socket = UnixSocket::bind().unwrap(); + let transport = Transport::Unix(socket.path().to_path_buf()); + let server = McpServer::new(FakeMcpServer::math()).serve_unix(socket); + let gateway = ready(gateway(vec![ServerSpec { name: "math".to_string(), transport }])).await; + assert_eq!(gateway.servers().len(), 1); + + server.stop(); + wait_until(&gateway, |gateway| matches!(gateway.statuses()[0].state, ServerState::Failed { .. })).await; + + assert!(gateway.servers().is_empty()); + assert!(!gateway.statuses()[0].can_sign_in); +} + +fn gateway(servers: Vec) -> McpGateway { + McpGateway::new(servers, &ClientOptions::default()) +} + +async fn ready(gateway: McpGateway) -> McpGateway { + gateway.ready().await; + gateway +} + +fn stdio(name: &str, command: &str, args: &[&str]) -> ServerSpec { + let args = args.iter().map(ToString::to_string).collect(); + let transport = Transport::Stdio { command: command.to_string(), args, env: HashMap::new(), cwd: None }; + ServerSpec { name: name.to_string(), transport } +} + +fn unreachable_http(name: &str) -> ServerSpec { + let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port(); + let url = format!("http://127.0.0.1:{port}/mcp"); + let transport = Transport::Http { url, headers: HeaderMap::new(), oauth: Some(OAuthSettings::default()) }; + ServerSpec { name: name.to_string(), transport } +} + +fn slow_server() -> FakeMcpServer { + FakeMcpServer::new().with_tool(FakeTool::new("slow").responds(FakeToolResponse::pending())) +} + +fn slow_call(server: &ConnectedServer) -> ToolCall { + server.call_tool("slow", Map::new(), ToolCallOptions::default()) +} + +async fn started(state: &FakeMcpState) { + state.wait_until(|state| !state.calls_for("slow").is_empty()).await; +} + +async fn refresh(state: &FakeMcpState) { + let requests = state.tool_list_requests(); + state.notify_tools_changed().await; + state.wait_until(|state| state.tool_list_requests() > requests).await; +} + +async fn fail_a_refresh(state: &FakeMcpState) { + let requests = state.tool_list_requests(); + state.fail_next_tool_list(); + state.clear_tools_and_notify().await; + state.wait_until(|state| state.tool_list_requests() > requests).await; +} + +fn prompt_text(prompt: &GetPromptResult) -> &str { + let ContentBlock::Text(text) = &prompt.messages[0].content else { panic!("expected text") }; + &text.text +} + +fn names(tools: &[Tool]) -> Vec { + tools.iter().map(|tool| tool.name.to_string()).collect() +} + +async fn released(alive: &Weak<()>) { + while alive.upgrade().is_some() { + tokio::task::yield_now().await; + } +} diff --git a/crates/mcp-utils/tests/http_headers.rs b/crates/mcp-utils/tests/http_headers.rs index f970566c5..1117bd831 100644 --- a/crates/mcp-utils/tests/http_headers.rs +++ b/crates/mcp-utils/tests/http_headers.rs @@ -1,137 +1,114 @@ use axum::{ Router, - extract::State, - http::{HeaderMap, HeaderName, Method, StatusCode, header::AUTHORIZATION}, - response::{IntoResponse, Response}, - routing::any, + extract::{Request, State}, + http::{HeaderMap, HeaderName, Method, header::AUTHORIZATION}, + middleware::{Next, from_fn_with_state}, + response::Response, }; -use mcp_utils::client::{McpConfig, McpHttpConfig, McpTransport}; -use rmcp::transport::streamable_http_client::{StreamableHttpClient, StreamableHttpClientTransportConfig}; +use mcp_utils::client::{ClientOptions, McpClient, Transport}; +use mcp_utils::model::{DiscoverResult, ErrorCode, ErrorData, ProtocolVersion, ServerCapabilities, ServerConfig}; +use reqwest::header::{HeaderMap as ConfiguredHeaders, HeaderValue}; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; +use rmcp::{RoleServer, ServerHandler, service::RequestContext}; +use std::future::{Future, ready}; +use std::sync::Arc; use tokio::{net::TcpListener, sync::mpsc, task::JoinHandle}; -use utils::variables::Vars; - -#[test] -fn configured_authorization_schemes_and_header_casing_are_preserved() { - for scheme in ["Sentry-Bearer", "Bearer", "bearer", "Basic", "Token"] { - for header_name in ["Authorization", "authorization", "AUTHORIZATION"] { - let config = parse_config( - "http://localhost/mcp", - &serde_json::json!({ - header_name: format!("{scheme} $TOKEN"), - "X-API-Key": "$API_KEY" - }), - ); - assert!(config.transport.auth_header.is_none()); - assert_eq!(config.transport.custom_headers[&AUTHORIZATION], format!("{scheme} secret")); - assert_eq!(config.transport.custom_headers[&HeaderName::from_static("x-api-key")], "custom-secret"); - } - } -} #[tokio::test] -async fn configured_headers_are_sent_verbatim_for_all_http_methods() { - let mut server = HeaderCaptureServer::start().await; - for method in [Method::POST, Method::GET, Method::DELETE] { - let transport = parse_config( - &server.url, - &serde_json::json!({ - "Authorization": "Sentry-Bearer $TOKEN", - "X-API-Key": "$API_KEY" - }), - ) - .transport; +async fn every_request_of_a_session_sends_the_configured_headers() { + let mut server = SessionServer::start().await; + let transport = Transport::Http { url: server.url.clone(), headers: configured_headers(), oauth: None }; + let client = McpClient::connect("remote", transport, &ClientOptions::default()).await.unwrap(); - send_request(method.clone(), transport).await; + let mut requests = server.requests_until(Method::GET).await; + client.close().await; + requests.extend(server.received()); - let (received_method, headers) = server.next_request().await; - assert_eq!(received_method, method); + for method in [Method::POST, Method::GET, Method::DELETE] { + let (_, headers) = + requests.iter().find(|(seen, _)| *seen == method).unwrap_or_else(|| panic!("the client sent no {method}")); assert_eq!(headers.get_all(AUTHORIZATION).iter().count(), 1, "{method}"); assert_eq!(headers[AUTHORIZATION], "Sentry-Bearer secret", "{method}"); assert_eq!(headers["x-api-key"], "custom-secret", "{method}"); } } -struct HeaderCaptureServer { +struct SessionServer { url: String, requests: mpsc::UnboundedReceiver<(Method, HeaderMap)>, task: JoinHandle<()>, } -impl HeaderCaptureServer { +impl SessionServer { async fn start() -> Self { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}/mcp", listener.local_addr().unwrap()); let (sender, requests) = mpsc::unbounded_channel(); - let app = Router::new().route("/mcp", any(capture_request)).with_state(sender); + let mcp = StreamableHttpService::new( + || Ok(LegacyServer), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + let app = Router::new().nest_service("/mcp", mcp).layer(from_fn_with_state(sender, record)); let task = tokio::spawn(async move { axum::serve(listener, app).await.unwrap(); }); Self { url, requests, task } } - async fn next_request(&mut self) -> (Method, HeaderMap) { - self.requests.recv().await.unwrap() + async fn requests_until(&mut self, last: Method) -> Vec<(Method, HeaderMap)> { + let mut requests = Vec::new(); + loop { + let request = self.requests.recv().await.unwrap(); + let done = request.0 == last; + requests.push(request); + if done { + return requests; + } + } + } + + fn received(&mut self) -> Vec<(Method, HeaderMap)> { + std::iter::from_fn(|| self.requests.try_recv().ok()).collect() } } -impl Drop for HeaderCaptureServer { +impl Drop for SessionServer { fn drop(&mut self) { self.task.abort(); } } -async fn capture_request( - State(sender): State>, - method: Method, - headers: HeaderMap, -) -> Response { - sender.send((method.clone(), headers)).unwrap(); - match method { - Method::POST => StatusCode::ACCEPTED.into_response(), - Method::GET => ([("content-type", "text/event-stream")], "").into_response(), - Method::DELETE => StatusCode::OK.into_response(), - _ => StatusCode::METHOD_NOT_ALLOWED.into_response(), +#[derive(Clone)] +struct LegacyServer; + +impl ServerHandler for LegacyServer { + fn get_info(&self) -> ServerConfig { + ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) + .with_protocol_version(ProtocolVersion::V_2025_11_25) + } + + fn discover( + &self, + _context: RequestContext, + ) -> impl Future> + Send + '_ { + ready(Err(ErrorData::new(ErrorCode::METHOD_NOT_FOUND, "server/discover is not supported", None))) } } -fn parse_config(url: &str, headers: &serde_json::Value) -> McpHttpConfig { - let json = serde_json::json!({"servers": {"remote": {"type": "http", "url": url, "headers": headers}}}); - let vars = Vars::new().with("TOKEN", "secret").with("API_KEY", "custom-secret"); - let mut servers = McpConfig::from_json(&json.to_string()).unwrap().into_servers(&vars).unwrap(); - let McpTransport::Http(config) = servers.remove(0).transport else { panic!("expected HTTP transport") }; - config +async fn record( + State(sender): State>, + request: Request, + next: Next, +) -> Response { + sender.send((request.method().clone(), request.headers().clone())).unwrap(); + next.run(request).await } -async fn send_request(method: Method, transport: StreamableHttpClientTransportConfig) { - let client = reqwest::Client::new(); - match method { - Method::POST => { - let message = - serde_json::from_value(serde_json::json!({"jsonrpc": "2.0", "method": "notifications/initialized"})) - .unwrap(); - client - .post_message(transport.uri, message, None, transport.auth_header, transport.custom_headers) - .await - .unwrap(); - } - Method::GET => { - let _stream = client - .get_stream( - transport.uri, - Some("session".into()), - None, - transport.auth_header, - transport.custom_headers, - ) - .await - .unwrap(); - } - Method::DELETE => { - client - .delete_session(transport.uri, "session".into(), transport.auth_header, transport.custom_headers) - .await - .unwrap(); - } - _ => panic!("unsupported test method: {method}"), - } +fn configured_headers() -> ConfiguredHeaders { + ConfiguredHeaders::from_iter([ + (AUTHORIZATION, HeaderValue::from_static("Sentry-Bearer secret")), + (HeaderName::from_static("x-api-key"), HeaderValue::from_static("custom-secret")), + ]) } diff --git a/crates/mcp-utils/tests/oauth.rs b/crates/mcp-utils/tests/oauth.rs new file mode 100644 index 000000000..d926321a8 --- /dev/null +++ b/crates/mcp-utils/tests/oauth.rs @@ -0,0 +1,294 @@ +use aether_auth::{OAuthClientRegistration, OAuthError}; +use mcp_utils::client::{ + ClientError, ClientOptions, Elicitation, ElicitationSink, McpClient, OAuthSettings, Transport, +}; +use mcp_utils::gateway::{GatewayError, McpGateway, ServerState}; +use mcp_utils::model::{ElicitResult, ElicitationAction, ElicitationCapability, FormElicitationCapability}; +use mcp_utils::testing::{ + ACCESS_TOKEN, CLIENT_METADATA_URL, FakeAuthServer, FakeMcpServer, OAuthPrompt, client_metadata, elicitation_host, + wait_until, +}; +use tokio::sync::mpsc; +use tokio::task::{JoinHandle, yield_now}; + +#[tokio::test] +async fn connect_reports_the_authorization_challenge() { + let server = math_auth_server().await; + let (options, _host) = elicitation_host(); + + let error = + McpClient::connect("linear", server.transport(), &options).await.expect_err("unauthorized connect fails"); + + assert!(matches!( + error, + ClientError::AuthRequired { server, challenge: Some(challenge) } + if server == "linear" && challenge.contains("resource_metadata") + )); +} + +#[tokio::test] +async fn servers_without_oauth_settings_are_not_offered_oauth() { + let server = math_auth_server().await; + let (options, _host) = elicitation_host(); + + let connected = McpClient::connect("linear", server.transport_with(None), &options).await; + let authorized = McpClient::authorize("linear", server.transport_with(None), &options, None).await; + + assert!(matches!(connected, Err(ClientError::Connect { server, .. }) if server == "linear")); + assert!(matches!(authorized, Err(ClientError::OAuthUnavailable { server }) if server == "linear")); +} + +#[tokio::test] +async fn oauth_is_not_offered_without_url_elicitation() { + let server = math_auth_server().await; + let (events, _host) = mpsc::channel(1); + let options = ClientOptions::default().elicitation(ElicitationSink::new( + events, + ElicitationCapability::new().with_form(FormElicitationCapability::new()), + )); + + let error = McpClient::connect("linear", server.transport(), &options).await.expect_err("connect fails"); + + assert!(matches!(error, ClientError::Connect { .. })); +} + +#[tokio::test] +async fn authorize_requires_an_event_sink() { + let server = math_auth_server().await; + + let error = McpClient::authorize("linear", server.transport(), &ClientOptions::default(), None) + .await + .expect_err("authorization needs a host to present the URL"); + + assert!(matches!(error, ClientError::OAuthUnavailable { server } if server == "linear")); +} + +#[tokio::test] +async fn authorization_presents_the_client_metadata_and_a_loopback_redirect() { + let server = math_auth_server().await; + let mut flow = AuthFlow::start(server.transport()); + + let prompt = flow.prompt().await; + let (client_id, redirect_uri) = (prompt.query("client_id"), prompt.query("redirect_uri")); + prompt.elicitation.respond(ElicitResult::new(ElicitationAction::Decline)); + + assert_eq!(client_id, CLIENT_METADATA_URL); + assert!(redirect_uri.starts_with("http://localhost:"), "{redirect_uri}"); + assert!(matches!( + flow.finish().await, + Err(ClientError::Auth { server, source: OAuthError::UserCancelled }) if server == "linear" + )); +} + +#[tokio::test] +async fn authorization_redirects_to_the_configured_callback_port() { + let server = math_auth_server().await; + let port = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port(); + let mut flow = + AuthFlow::start(server.transport_with(Some(OAuthSettings { callback_port: port, ..client_metadata() }))); + + let prompt = flow.prompt().await; + let redirect_uri = prompt.query("redirect_uri"); + drop(prompt); + + assert_eq!(redirect_uri, format!("http://localhost:{port}/")); + assert!(matches!(flow.finish().await, Err(ClientError::Auth { source: OAuthError::UserCancelled, .. }))); +} + +#[tokio::test] +async fn authorization_presents_a_preregistered_client_id() { + let server = math_auth_server().await; + let oauth = OAuthSettings { + registration: OAuthClientRegistration::PreRegistered("registered".to_string()), + callback_port: 0, + }; + let mut flow = AuthFlow::start(server.transport_with(Some(oauth))); + + let prompt = flow.prompt().await; + let client_id = prompt.query("client_id"); + drop(prompt); + + assert_eq!(client_id, "registered"); + assert!(matches!(flow.finish().await, Err(ClientError::Auth { source: OAuthError::UserCancelled, .. }))); +} + +#[tokio::test] +async fn accepted_authorization_waits_for_the_callback_then_connects_with_the_token() { + let server = math_auth_server().await; + let mut flow = AuthFlow::start(server.transport()); + + let prompt = flow.prompt().await; + let callback = prompt.callback_url(); + prompt.elicitation.respond(ElicitResult::new(ElicitationAction::Accept)); + yield_now().await; + assert!(!flow.task.is_finished(), "accepting the prompt does not complete the browser flow"); + reqwest::get(callback).await.expect("loopback callback responds"); + + assert!(matches!( + flow.host.recv().await, + Some(Elicitation::Complete { server, id }) if server == "linear" && id == "oauth" + )); + let client = flow.finish().await.expect("the token authorizes the MCP handshake"); + assert!(client.list_tools().await.unwrap().iter().any(|tool| tool.name == "add_numbers")); + assert!(server.received_bearer(ACCESS_TOKEN)); +} + +#[tokio::test] +async fn gateway_servers_that_demand_oauth_need_it() { + let server = math_auth_server().await; + let (gateway, _host) = oauth_gateway(&server); + + gateway.ready().await; + + let status = &gateway.statuses()[0]; + assert_eq!(status.state, ServerState::NeedsOAuth); + assert!(status.can_sign_in); + assert!(gateway.servers().is_empty()); +} + +#[tokio::test] +async fn gateway_servers_that_demand_oauth_fail_without_url_elicitation() { + let server = math_auth_server().await; + let (events, _host) = mpsc::channel(1); + let options = ClientOptions::default().elicitation(ElicitationSink::new( + events, + ElicitationCapability::new().with_form(FormElicitationCapability::new()), + )); + let gateway = McpGateway::new(vec![server.spec("linear")], &options); + + gateway.ready().await; + let status = &gateway.statuses()[0]; + assert!(matches!(status.state, ServerState::Failed { .. })); + assert!(!status.can_sign_in); +} + +#[tokio::test] +async fn gateway_authentication_reconnects_the_server_once_authorized() { + let server = math_auth_server().await; + let (gateway, mut host) = oauth_gateway(&server); + gateway.ready().await; + + let authentication = authenticate(&gateway, "linear"); + let prompt = next_prompt(&mut host).await; + let callback = prompt.callback_url(); + prompt.elicitation.respond(ElicitResult::new(ElicitationAction::Accept)); + reqwest::get(callback).await.expect("loopback callback responds"); + + authentication.await.unwrap().expect("authentication succeeds"); + assert_eq!(gateway.statuses()[0].state, ServerState::Connected { tool_count: 2 }); + assert!(gateway.statuses()[0].can_sign_in); + assert!(gateway.servers()[0].tools().iter().any(|tool| tool.name == "add_numbers")); + assert!(matches!(host.recv().await, Some(Elicitation::Complete { server, .. }) if server == "linear")); +} + +#[tokio::test] +async fn gateway_authentication_reports_a_declined_authorization_and_allows_a_retry() { + let server = math_auth_server().await; + let (gateway, mut host) = oauth_gateway(&server); + gateway.ready().await; + + let authentication = authenticate(&gateway, "linear"); + wait_until(&gateway, |gateway| gateway.statuses()[0].state == ServerState::Authenticating).await; + next_prompt(&mut host).await.elicitation.respond(ElicitResult::new(ElicitationAction::Decline)); + + let declined = authentication.await.unwrap(); + assert!( + matches!(declined, Err(GatewayError::Client(ClientError::Auth { server, source: OAuthError::UserCancelled })) if server == "linear") + ); + let status = gateway.statuses().remove(0); + assert!(matches!(status.state, ServerState::Failed { .. })); + assert!(status.can_sign_in); + let retry = authenticate(&gateway, "linear"); + drop(next_prompt(&mut host).await); + assert!(matches!( + retry.await.unwrap(), + Err(GatewayError::Client(ClientError::Auth { source: OAuthError::UserCancelled, .. })) + )); +} + +#[tokio::test] +async fn authenticating_again_cancels_the_pending_authentication() { + let server = math_auth_server().await; + let (gateway, mut host) = oauth_gateway(&server); + gateway.ready().await; + let authentication = authenticate(&gateway, "linear"); + let _prompt = next_prompt(&mut host).await; + + let _again = authenticate(&gateway, "linear"); + + let cancelled = authentication.await.unwrap(); + assert!( + matches!(cancelled, Err(GatewayError::Client(ClientError::Auth { server, source: OAuthError::UserCancelled })) if server == "linear") + ); +} + +#[tokio::test] +async fn shutting_down_abandons_a_pending_authentication() { + let server = math_auth_server().await; + let (gateway, mut host) = oauth_gateway(&server); + gateway.ready().await; + let authentication = authenticate(&gateway, "linear"); + let _prompt = next_prompt(&mut host).await; + + gateway.shutdown().await; + + assert!(matches!(authentication.await.unwrap(), Err(GatewayError::Closed))); +} + +#[tokio::test] +async fn a_configured_callback_port_that_is_taken_fails_authorization() { + let server = math_auth_server().await; + let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let oauth = OAuthSettings { callback_port: taken.local_addr().unwrap().port(), ..client_metadata() }; + let (options, _host) = elicitation_host(); + + let error = McpClient::authorize("linear", server.transport_with(Some(oauth)), &options, None) + .await + .expect_err("the callback port is in use"); + + assert!(matches!( + error, + ClientError::Auth { source: OAuthError::Io(error), .. } if error.kind() == std::io::ErrorKind::AddrInUse + )); +} + +struct AuthFlow { + task: JoinHandle>, + host: mpsc::Receiver, +} + +impl AuthFlow { + fn start(transport: Transport) -> Self { + let (options, host) = elicitation_host(); + let task = tokio::spawn(async move { McpClient::authorize("linear", transport, &options, None).await }); + Self { task, host } + } + + async fn prompt(&mut self) -> OAuthPrompt { + next_prompt(&mut self.host).await + } + + async fn finish(self) -> Result { + self.task.await.unwrap() + } +} + +fn oauth_gateway(server: &FakeAuthServer) -> (McpGateway, mpsc::Receiver) { + let (options, host) = elicitation_host(); + (McpGateway::new(vec![server.spec("linear")], &options), host) +} + +fn authenticate(gateway: &McpGateway, server: &str) -> JoinHandle> { + let (gateway, server) = (gateway.clone(), server.to_string()); + tokio::spawn(async move { gateway.authenticate(&server).await }) +} + +async fn next_prompt(host: &mut mpsc::Receiver) -> OAuthPrompt { + let prompt = OAuthPrompt::receive(host).await; + assert_eq!(prompt.elicitation.server, "linear"); + prompt +} + +async fn math_auth_server() -> FakeAuthServer { + FakeAuthServer::start(FakeMcpServer::math()).await +} diff --git a/crates/mcp-utils/tests/oauth_config.rs b/crates/mcp-utils/tests/oauth_config.rs deleted file mode 100644 index bf5dd929e..000000000 --- a/crates/mcp-utils/tests/oauth_config.rs +++ /dev/null @@ -1,61 +0,0 @@ -use aether_auth::OAuthClientRegistration; -use mcp_utils::client::{McpConfig, McpHttpConfig, McpTransport}; -use utils::variables::Vars; - -#[tokio::test] -async fn http_oauth_defaults_to_aether_cimd_and_fixed_callback() { - let config = parse_http_config(r#"{ "type": "http", "url": "https://example.com/mcp" }"#, &Vars::new()); - let oauth = config.resolved_oauth().expect("OAuth should be offered without an auth header"); - - assert_eq!( - oauth.client_registration, - OAuthClientRegistration::ClientMetadata("https://aether-agent.io/oauth/client-metadata.json".to_string()) - ); - assert_eq!(oauth.callback_port.get(), 3118); - assert_eq!(oauth.redirect_uri(), "http://localhost:3118/"); -} - -#[tokio::test] -async fn custom_cimd_expands_variables_and_preregistered_client_takes_priority() { - let vars = Vars::new().with("CLIENT_ID", "registered").with("CIMD_URL", "https://client.example/oauth/client.json"); - let config = parse_http_config( - r#"{ - "type": "http", - "url": "https://example.com/mcp", - "oauth": { - "clientId": "$CLIENT_ID", - "clientMetadataUrl": "$CIMD_URL", - "callbackPort": 4000 - } - }"#, - &vars, - ); - let oauth = config.resolved_oauth().expect("OAuth should be offered without an auth header"); - - assert_eq!(oauth.client_registration, OAuthClientRegistration::PreRegistered("registered".to_string())); - assert_eq!(oauth.callback_port.get(), 4000); - assert_eq!(config.oauth.unwrap().client_metadata_url.as_deref(), Some("https://client.example/oauth/client.json")); -} - -#[tokio::test] -async fn bearer_header_bypasses_oauth_defaults() { - let config = parse_http_config( - r#"{ - "type": "http", - "url": "https://example.com/mcp", - "headers": { "Authorization": "Bearer token" } - }"#, - &Vars::new(), - ); - - assert!(config.resolved_oauth().is_none()); -} - -fn parse_http_config(server: &str, vars: &Vars) -> McpHttpConfig { - let json = format!(r#"{{ "servers": {{ "remote": {server} }} }}"#); - let mut servers = McpConfig::from_json(&json).unwrap().into_servers(vars).unwrap(); - assert_eq!(servers.len(), 1); - let server = servers.remove(0); - let McpTransport::Http(config) = server.transport else { panic!("expected HTTP transport") }; - config -} diff --git a/crates/mcp-utils/tests/server.rs b/crates/mcp-utils/tests/server.rs new file mode 100644 index 000000000..d384cbce0 --- /dev/null +++ b/crates/mcp-utils/tests/server.rs @@ -0,0 +1,98 @@ +use mcp_utils::client::{ClientOptions, McpClient, Transport}; +use mcp_utils::server::{McpServer, ServerHandle, UnixSocket}; +use mcp_utils::testing::{FakeMcpServer, FakeTool, FakeToolResponse, call}; +use serde_json::json; +use std::fs::metadata; +use std::os::unix::fs::PermissionsExt; + +#[tokio::test] +async fn serve_unix_binds_a_socket_in_a_private_directory() { + let handle = McpServer::new(FakeMcpServer::new()).serve_unix(UnixSocket::bind().unwrap()); + + let directory = handle.path().parent().unwrap(); + assert!(handle.path().exists()); + assert_eq!(metadata(directory).unwrap().permissions().mode() & 0o777, 0o700); +} + +#[tokio::test] +async fn a_socket_is_bound_before_a_server_is_attached_to_it() { + let socket = UnixSocket::bind().unwrap(); + let path = socket.path().to_path_buf(); + assert!(path.exists()); + + let handle = McpServer::new(FakeMcpServer::math()).serve_unix(socket); + let tools = connect(&handle).await.list_tools().await.unwrap(); + + assert_eq!(handle.path(), path); + assert_eq!(tools.len(), 2); +} + +#[tokio::test] +async fn dropping_the_handle_removes_the_socket_and_its_directory() { + let handle = McpServer::new(FakeMcpServer::new()).serve_unix(UnixSocket::bind().unwrap()); + let socket = handle.path().to_path_buf(); + let directory = socket.parent().unwrap().to_path_buf(); + + drop(handle); + + assert!(!socket.exists()); + assert!(!directory.exists()); +} + +#[tokio::test] +async fn stopping_the_handle_removes_the_socket_and_ends_open_connections() { + let handle = McpServer::new(FakeMcpServer::new()).serve_unix(UnixSocket::bind().unwrap()); + let client = connect(&handle).await; + + handle.stop(); + + assert!(!handle.path().exists()); + assert!(!handle.path().parent().unwrap().exists()); + assert!(client.list_tools().await.is_err()); +} + +#[tokio::test] +async fn unix_clients_reach_the_same_server() { + let server = FakeMcpServer::math(); + let state = server.state(); + let handle = McpServer::new(server).serve_unix(UnixSocket::bind().unwrap()); + let first = connect(&handle).await; + let second = connect(&handle).await; + + call(&first, "add_numbers", json!({"a": 1, "b": 2})).await.unwrap(); + call(&second, "add_numbers", json!({"a": 3, "b": 4})).await.unwrap(); + + assert_eq!(state.calls_for("add_numbers").len(), 2); +} + +#[tokio::test] +async fn dropping_the_handle_cancels_in_flight_calls() { + let server = FakeMcpServer::new().with_tool(FakeTool::new("slow").responds(FakeToolResponse::pending())); + let state = server.state(); + let handle = McpServer::new(server).serve_unix(UnixSocket::bind().unwrap()); + let client = connect(&handle).await; + let in_flight = tokio::spawn(async move { call(&client, "slow", json!({})).await }); + state.wait_until(|state| !state.calls_for("slow").is_empty()).await; + + drop(handle); + + assert!(in_flight.await.unwrap().is_err()); +} + +#[tokio::test] +async fn closing_an_in_memory_client_stops_its_server() { + let server = FakeMcpServer::new(); + let alive = server.alive(); + let client = + McpClient::connect("test", Transport::InProcess(server.into()), &ClientOptions::default()).await.unwrap(); + + client.close().await; + + while alive.upgrade().is_some() { + tokio::task::yield_now().await; + } +} + +async fn connect(handle: &ServerHandle) -> McpClient { + McpClient::connect("test", Transport::Unix(handle.path().to_path_buf()), &ClientOptions::default()).await.unwrap() +} diff --git a/crates/mcp-utils/tests/tool_catalog.rs b/crates/mcp-utils/tests/tool_catalog.rs deleted file mode 100644 index a5af6499f..000000000 --- a/crates/mcp-utils/tests/tool_catalog.rs +++ /dev/null @@ -1,147 +0,0 @@ -use mcp_utils::client::{ - CatalogTool, DeferredToolRules, PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME, ServerCatalogEntry, ToolAnnotationMatcher, - ToolCatalog, ToolExposure, ToolExposureKind, ToolFilter, ToolMatcher, ToolRoute, -}; -use rmcp::model::{Tool, ToolAnnotations}; -use serde_json::json; -use std::sync::Arc; -use utils::mcp_status::{McpServerAuthCapability, McpServerStatus}; - -fn tool(name: &str) -> Tool { - let schema = json!({ - "type": "object", - "properties": { "value": { "type": "string" } }, - "required": ["value"] - }); - Tool::new(name.to_string(), format!("{name} description"), Arc::new(schema.as_object().unwrap().clone())) -} - -fn connected_entry(name: &str, exposure: ToolExposure, tools: &[Tool], filter: &ToolFilter) -> ServerCatalogEntry { - ServerCatalogEntry::new( - name, - format!("{name} server"), - Some(format!("{name} instructions")), - McpServerStatus::Connected { tool_count: tools.len() }, - McpServerAuthCapability::OAuth, - exposure, - tools, - filter, - ) -} - -#[test] -fn catalog_projects_visibility_filtering_instructions_and_routes_consistently() { - assert_eq!(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME, "progressive-discovery"); - let tools = vec![ - tool("bash"), - tool("read").with_annotations(ToolAnnotations::new().read_only(true).open_world(false)), - tool("write"), - ]; - let filter = ToolFilter { - allow: vec![ToolMatcher::name("coding__*")], - deny: vec![ToolMatcher::annotations(ToolAnnotationMatcher { - read_only: Some(true), - ..ToolAnnotationMatcher::default() - })], - }; - let exposure = ToolExposure::Deferred(DeferredToolRules::new(&["read", "write"], &["bash"])); - let mut catalog = ToolCatalog::new(); - catalog.upsert_server(connected_entry("coding", exposure, &tools, &filter)); - catalog.set_progressive_discovery_instructions(Some("Discover deferred tools".to_string())); - - let tools = catalog.tools(); - assert_eq!( - tools.model_visible.iter().map(|tool| tool.definition().name.as_str()).collect::>(), - ["coding__bash"] - ); - assert_eq!(tools.deferred.into_iter().map(CatalogTool::local_name).collect::>(), ["write"]); - let server_tools = catalog.tools_for("coding").unwrap(); - assert_eq!(server_tools.model_visible.len(), 1); - assert_eq!(server_tools.deferred.len(), 1); - assert!(catalog.tools_for("missing").is_none()); - assert_eq!(catalog.discoverable_deferred_servers()[0].name, "coding"); - assert_eq!(catalog.model_instructions().get("coding").map(String::as_str), Some("coding instructions")); - assert_eq!( - catalog.model_instructions().get(PROGRESSIVE_DISCOVERY_INSTRUCTION_NAME).map(String::as_str), - Some("Discover deferred tools") - ); - assert!(catalog.route_permitted(&ToolRoute::ModelVisible { namespaced_name: "coding__bash".into() })); - assert!(!catalog.route_permitted(&ToolRoute::Deferred { server: "coding".into(), tool: "bash".into() })); - assert!(catalog.route_permitted(&ToolRoute::Deferred { server: "coding".into(), tool: "write".into() })); - assert!(!catalog.route_permitted(&ToolRoute::Deferred { server: "coding".into(), tool: "read".into() })); -} - -#[test] -fn catalog_preserves_names_schema_annotations_metadata_and_order() { - let tools = - vec![tool("first").with_annotations(ToolAnnotations::new().read_only(true).idempotent(true)), tool("second")]; - let mut catalog = ToolCatalog::new(); - catalog.upsert_server(connected_entry("alpha", ToolExposure::ModelVisible, &tools, &ToolFilter::default())); - catalog.upsert_server(connected_entry("beta", ToolExposure::deferred_all(), &tools, &ToolFilter::default())); - - let first = catalog.tool("alpha__first").unwrap(); - assert_eq!(first.namespaced_name(), "alpha__first"); - assert_eq!(first.local_name(), "first"); - assert_eq!(first.definition().server.as_deref(), Some("alpha")); - assert_eq!(first.definition().description, "first description"); - assert_eq!(first.definition().parameters["required"], json!(["value"])); - let annotations = first.definition().annotations.as_ref().unwrap(); - assert_eq!(annotations.read_only_hint, Some(true)); - assert_eq!(annotations.idempotent_hint, Some(true)); - assert_eq!(first.exposure(), ToolExposureKind::ModelVisible); - assert!(first.allowed()); - assert_eq!(catalog.servers().iter().map(ServerCatalogEntry::name).collect::>(), ["alpha", "beta"]); - assert_eq!( - catalog.servers()[0].tools().iter().map(CatalogTool::local_name).collect::>(), - ["first", "second"] - ); - assert_eq!(catalog.servers()[0].description(), "alpha server"); - assert_eq!(catalog.servers()[0].auth_capability(), McpServerAuthCapability::OAuth); -} - -#[test] -fn disconnected_servers_retain_status_without_exposing_tools_or_instructions() { - let tools = vec![tool("hidden")]; - let mut catalog = ToolCatalog::new(); - catalog.upsert_server(ServerCatalogEntry::new( - "remote", - "Remote", - Some("not visible".to_string()), - McpServerStatus::NeedsOAuth, - McpServerAuthCapability::OAuth, - ToolExposure::deferred_all(), - &tools, - &ToolFilter::default(), - )); - catalog.set_progressive_discovery_instructions(Some("not visible".to_string())); - - let tools = catalog.tools(); - assert_eq!(tools.model_visible.len(), 0); - assert_eq!(tools.deferred.len(), 0); - assert_eq!(catalog.discoverable_deferred_servers().len(), 0); - assert!(catalog.model_instructions().is_empty()); - assert!(!catalog.route_permitted(&ToolRoute::Deferred { server: "remote".into(), tool: "hidden".into() })); - assert!(matches!(catalog.servers()[0].status(), McpServerStatus::NeedsOAuth)); - assert!(catalog.servers()[0].status_entry().can_authenticate()); -} - -#[test] -fn replacement_preserves_server_position_and_removal_is_atomic() { - let mut catalog = ToolCatalog::new(); - catalog.upsert_server(connected_entry("first", ToolExposure::ModelVisible, &[tool("old")], &ToolFilter::default())); - catalog.upsert_server(connected_entry( - "second", - ToolExposure::ModelVisible, - &[tool("other")], - &ToolFilter::default(), - )); - catalog.upsert_server(connected_entry("first", ToolExposure::ModelVisible, &[tool("new")], &ToolFilter::default())); - - assert_eq!(catalog.servers().iter().map(ServerCatalogEntry::name).collect::>(), ["first", "second"]); - assert!(catalog.tool("first__old").is_none()); - assert!(catalog.tool("first__new").is_some()); - - let removed = catalog.remove_server("first").unwrap(); - assert_eq!(removed.name(), "first"); - assert_eq!(catalog.servers().iter().map(ServerCatalogEntry::name).collect::>(), ["second"]); -} diff --git a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx index 3ca4453ab..a2f7b0639 100644 --- a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx +++ b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx @@ -128,37 +128,11 @@ If a server entry has a `command` and no `type`, it is treated as `"type": "stdi | `deferTools` | `boolean \| object` | `true` proxies every tool; an object partitions tools with `include` and `exclude` patterns. | - - Remote MCP servers over HTTP with Server-Sent Events. Internally, Aether uses the same HTTP transport path. - - ```json title=".aether/mcp.json" - { - "servers": { - "remote-tools": { - "type": "sse", - "url": "https://mcp.example.com/sse", - "headers": { - "Authorization": "Bearer $MCP_TOKEN" - } - } - } - } - ``` - - | Field | Type | Description | - |-------|------|-------------| - | `type` | `"sse"` | Required. | - | `url` | `string` | SSE endpoint URL. | - | `headers` | `object` | Optional headers. Aether currently forwards the `Authorization` header to the transport. | - | `oauth` | `object` | Optional custom CIMD or pre-registered public OAuth client settings. See [OAuth clients](#oauth-clients). | - | `deferTools` | `boolean \| object` | `true` proxies every tool; an object partitions tools with `include` and `exclude` patterns. | - - ## OAuth clients -For `http` and `sse` servers that return an OAuth challenge, Aether defaults to its first-party Client ID Metadata Document (CIMD) at `https://aether-agent.io/oauth/client-metadata.json`. The callback listener binds `http://localhost:3118/`. If port 3118 is occupied, authentication fails rather than selecting an unregistered redirect URI. +For `http` servers that return an OAuth challenge, Aether defaults to its first-party Client ID Metadata Document (CIMD) at `https://aether-agent.io/oauth/client-metadata.json`. The callback listener binds `http://localhost:3118/`. If port 3118 is occupied, authentication fails rather than selecting an unregistered redirect URI. Set `oauth.clientMetadataUrl` to use a custom CIMD, or `oauth.clientId` for a public client pre-registered with the provider. A pre-registered client ID takes priority when both are configured: @@ -350,7 +324,7 @@ Only the private session socket and Aether executable directory are injected int ## Variable expansion -Variables are expanded in two places: MCP source paths in settings, and server fields inside MCP JSON files. Environment variables are expanded in stdio `command`, stdio `args`, stdio `env` values, in-memory `args`, HTTP/SSE `url`, and the HTTP/SSE `Authorization` header value. `${WORKSPACE}` and `${AETHER_HOME}` (the user-level Aether home, honoring the `AETHER_HOME` env var) are also available when server configs are converted into running MCP servers. +Variables are expanded in two places: MCP source paths in settings, and server fields inside MCP JSON files. Environment variables are expanded in stdio `command`, stdio `args`, stdio `env` values, in-memory `args`, HTTP `url`, and HTTP header values. `${WORKSPACE}` and `${AETHER_HOME}` (the user-level Aether home, honoring the `AETHER_HOME` env var) are also available when server configs are converted into running MCP servers. Supported syntax: