From 8aabc7112d5640efa9df2dcc3e39875ff11783fa Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:13 -0700 Subject: [PATCH 1/6] refactor(mcp-utils): separate MCP clients, gateways, and servers Replace the client manager and snapshot architecture with explicit McpClient, McpGateway, and McpServer APIs so transport and connection lifecycle are independent of agent-specific policy. Consolidate transports, OAuth, elicitation, tool calls, and task handling behind the new APIs. Move configuration, tool filtering, and model-facing discovery out of this crate. Replace legacy test helpers with in-memory server and task fakes, and cover the client, gateway, server, and OAuth public APIs. --- .aether/skills/rust/error-handling.md | 2 - crates/mcp-utils/Cargo.toml | 52 +- crates/mcp-utils/README.md | 8 +- crates/mcp-utils/src/client/call_tool.rs | 221 ---- crates/mcp-utils/src/client/config.rs | 804 ------------ crates/mcp-utils/src/client/connection.rs | 385 ------ .../src/client/connection_attempt_manager.rs | 39 - crates/mcp-utils/src/client/elicitation.rs | 76 +- crates/mcp-utils/src/client/error.rs | 119 +- crates/mcp-utils/src/client/handler.rs | 136 +++ crates/mcp-utils/src/client/manager.rs | 1086 ----------------- crates/mcp-utils/src/client/mcp_client.rs | 332 ++--- crates/mcp-utils/src/client/mcp_snapshot.rs | 100 -- crates/mcp-utils/src/client/mod.rs | 53 +- crates/mcp-utils/src/client/mrtr.rs | 285 ----- crates/mcp-utils/src/client/naming.rs | 9 - crates/mcp-utils/src/client/oauth.rs | 137 +++ crates/mcp-utils/src/client/oauth_handler.rs | 152 --- crates/mcp-utils/src/client/options.rs | 69 ++ crates/mcp-utils/src/client/task.rs | 529 ++------ crates/mcp-utils/src/client/tool_call.rs | 329 +++++ crates/mcp-utils/src/client/tool_catalog.rs | 295 ----- crates/mcp-utils/src/client/tool_filter.rs | 85 -- crates/mcp-utils/src/client/transport.rs | 176 +++ .../mcp-utils/src/docs/mcp_server_config.md | 64 - crates/mcp-utils/src/gateway/connection.rs | 225 ++++ crates/mcp-utils/src/gateway/error.rs | 12 + crates/mcp-utils/src/gateway/mcp_gateway.rs | 172 +++ crates/mcp-utils/src/gateway/mod.rs | 6 + crates/mcp-utils/src/lib.rs | 12 +- crates/mcp-utils/src/protocol.rs | 14 - crates/mcp-utils/src/server.rs | 194 +++ crates/mcp-utils/src/server/mod.rs | 2 - crates/mcp-utils/src/server/mrtr.rs | 123 -- crates/mcp-utils/src/server/tasks.rs | 15 - crates/mcp-utils/src/server/tool.rs | 0 crates/mcp-utils/src/testing.rs | 264 +--- .../src/testing/elicitation_script.rs | 106 ++ .../mcp-utils/src/testing/fake_auth_server.rs | 172 +++ crates/mcp-utils/src/testing/fake_mcp.rs | 388 ++++-- crates/mcp-utils/src/testing/fake_task.rs | 138 +++ crates/mcp-utils/src/testing/gate.rs | 40 + crates/mcp-utils/src/tool_gateway/mod.rs | 9 - .../mcp-utils/src/tool_gateway/transport.rs | 294 ----- crates/mcp-utils/tests/client.rs | 788 ++++++++++++ crates/mcp-utils/tests/doc_examples.rs | 10 - .../tests/elicitation_notifications.rs | 38 - crates/mcp-utils/tests/gateway.rs | 327 +++++ crates/mcp-utils/tests/http_headers.rs | 165 ++- crates/mcp-utils/tests/oauth.rs | 294 +++++ crates/mcp-utils/tests/oauth_config.rs | 61 - crates/mcp-utils/tests/server.rs | 98 ++ crates/mcp-utils/tests/tool_catalog.rs | 147 --- 53 files changed, 4155 insertions(+), 5502 deletions(-) delete mode 100644 crates/mcp-utils/src/client/call_tool.rs delete mode 100644 crates/mcp-utils/src/client/config.rs delete mode 100644 crates/mcp-utils/src/client/connection.rs delete mode 100644 crates/mcp-utils/src/client/connection_attempt_manager.rs create mode 100644 crates/mcp-utils/src/client/handler.rs delete mode 100644 crates/mcp-utils/src/client/manager.rs delete mode 100644 crates/mcp-utils/src/client/mcp_snapshot.rs delete mode 100644 crates/mcp-utils/src/client/mrtr.rs delete mode 100644 crates/mcp-utils/src/client/naming.rs create mode 100644 crates/mcp-utils/src/client/oauth.rs delete mode 100644 crates/mcp-utils/src/client/oauth_handler.rs create mode 100644 crates/mcp-utils/src/client/options.rs create mode 100644 crates/mcp-utils/src/client/tool_call.rs delete mode 100644 crates/mcp-utils/src/client/tool_catalog.rs delete mode 100644 crates/mcp-utils/src/client/tool_filter.rs create mode 100644 crates/mcp-utils/src/client/transport.rs delete mode 100644 crates/mcp-utils/src/docs/mcp_server_config.md create mode 100644 crates/mcp-utils/src/gateway/connection.rs create mode 100644 crates/mcp-utils/src/gateway/error.rs create mode 100644 crates/mcp-utils/src/gateway/mcp_gateway.rs create mode 100644 crates/mcp-utils/src/gateway/mod.rs delete mode 100644 crates/mcp-utils/src/protocol.rs create mode 100644 crates/mcp-utils/src/server.rs delete mode 100644 crates/mcp-utils/src/server/mod.rs delete mode 100644 crates/mcp-utils/src/server/mrtr.rs delete mode 100644 crates/mcp-utils/src/server/tasks.rs delete mode 100644 crates/mcp-utils/src/server/tool.rs create mode 100644 crates/mcp-utils/src/testing/elicitation_script.rs create mode 100644 crates/mcp-utils/src/testing/fake_auth_server.rs create mode 100644 crates/mcp-utils/src/testing/fake_task.rs create mode 100644 crates/mcp-utils/src/testing/gate.rs delete mode 100644 crates/mcp-utils/src/tool_gateway/mod.rs delete mode 100644 crates/mcp-utils/src/tool_gateway/transport.rs create mode 100644 crates/mcp-utils/tests/client.rs delete mode 100644 crates/mcp-utils/tests/doc_examples.rs delete mode 100644 crates/mcp-utils/tests/elicitation_notifications.rs create mode 100644 crates/mcp-utils/tests/gateway.rs create mode 100644 crates/mcp-utils/tests/oauth.rs delete mode 100644 crates/mcp-utils/tests/oauth_config.rs create mode 100644 crates/mcp-utils/tests/server.rs delete mode 100644 crates/mcp-utils/tests/tool_catalog.rs 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/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/tool_filter.rs b/crates/mcp-utils/src/client/tool_filter.rs deleted file mode 100644 index 070ef3e5f..000000000 --- a/crates/mcp-utils/src/client/tool_filter.rs +++ /dev/null @@ -1,85 +0,0 @@ -use llm::ToolDefinition; -use utils::matches_name_pattern; - -#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] -#[serde(untagged)] -pub enum ToolMatcher { - Name(String), - Annotations(ToolAnnotationMatcher), -} - -impl ToolMatcher { - pub fn name(pattern: impl Into) -> Self { - Self::Name(pattern.into()) - } - - pub fn read_only() -> Self { - Self::Annotations(ToolAnnotationMatcher { read_only: Some(true), ..ToolAnnotationMatcher::default() }) - } - - pub fn annotations(matcher: ToolAnnotationMatcher) -> Self { - Self::Annotations(matcher) - } - - pub fn matches(&self, tool: &ToolDefinition) -> bool { - match self { - Self::Name(pattern) => matches_name_pattern(pattern, &tool.name), - Self::Annotations(matcher) => matcher.matches(tool), - } - } -} - -#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub struct ToolAnnotationMatcher { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub read_only: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub destructive: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub idempotent: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub open_world: Option, -} - -impl ToolAnnotationMatcher { - pub fn matches(&self, tool: &ToolDefinition) -> bool { - let Some(annotations) = tool.annotations.as_ref() else { return false }; - let pairs = [ - (self.read_only, annotations.read_only_hint), - (self.destructive, annotations.destructive_hint), - (self.idempotent, annotations.idempotent_hint), - (self.open_world, annotations.open_world_hint), - ]; - if pairs.iter().all(|(field, _)| field.is_none()) { - return false; - } - pairs.iter().all(|(field, hint)| field.is_none_or(|value| *hint == Some(value))) - } -} - -/// 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 { - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub allow: Vec, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub deny: Vec, -} - -impl ToolFilter { - pub fn is_empty(&self) -> bool { - 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)); - allowed && !denied - } -} 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/docs/mcp_server_config.md b/crates/mcp-utils/src/docs/mcp_server_config.md deleted file mode 100644 index 015b170c1..000000000 --- a/crates/mcp-utils/src/docs/mcp_server_config.md +++ /dev/null @@ -1,64 +0,0 @@ -A single MCP server. The `type` field selects the transport; `stdio` is the -default and may be omitted. - -A local stdio server launched as a subprocess: - -```json -{ - "type": "stdio", - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-github"], - "env": { "GITHUB_TOKEN": "ghp_..." } -} -``` - -A remote streamable-HTTP server using a pre-registered public OAuth client: - -```json -{ - "type": "http", - "url": "https://mcp.slack.com/mcp", - "oauth": { - "clientId": "1601185624273.8899143856786", - "callbackPort": 3118 - } -} -``` - -When a remote HTTP server returns an OAuth challenge, Aether uses its first-party -Client ID Metadata Document at -`https://aether-agent.io/oauth/client-metadata.json` and listens on -`127.0.0.1:3118`, advertising the exact redirect URI `http://localhost:3118/`. -If port 3118 is occupied, authentication fails until the port is available. - -The `oauth` object is optional. Set `clientMetadataUrl` for a custom CIMD, or -`clientId` for a pre-registered public client; a configured `clientId` takes -priority. `callbackPort` defaults to 3118 and must exactly match the client's -registered redirect URI. If the authorization server does not advertise CIMD, -Aether falls back to deprecated Dynamic Client Registration. An explicit -`Authorization` header bypasses OAuth entirely. - -A remote server using a bearer token: - -```json -{ - "type": "http", - "url": "https://mcp.example.com", - "headers": { "Authorization": "Bearer ..." } -} -``` - -Set `"deferTools": true` to omit every tool on this server from the model-visible -tool definitions and make it available through progressive `aether mcp` discovery. -For selective deferral, set `deferTools` to an object with `include` and `exclude` -lists. Entries match either an exact MCP-local tool name or a prefix ending in `*`: - -```json -{ - "type": "in-memory", - "deferTools": { - "include": ["*"], - "exclude": ["bash", "lsp_*"] - } -} -``` 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"]); -} From d7a4258253765c5c20c7f29ae8329f3d37e6a779 Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:13 -0700 Subject: [PATCH 2/6] refactor(aether-core): own MCP configuration and tool discovery Introduce McpCatalog as the agent-facing view over McpGateway, with configuration, tool filtering, and direct or deferred exposure managed by aether-core. Move progressive discovery instructions and deferred-tool IPC into the core runtime. Synchronize agent tools and server instructions directly from catalog changes instead of external update commands. Replace the MCP handle and runtime abstractions, adapt tool and task execution, and migrate examples and tests to the new public APIs. --- Cargo.lock | 1 + README.md | 10 +- crates/aether-core/Cargo.toml | 2 + crates/aether-core/README.md | 12 +- crates/aether-core/examples/mcp_agent.rs | 8 +- crates/aether-core/src/agent_spec.rs | 232 +----- crates/aether-core/src/core/agent.rs | 109 +-- crates/aether-core/src/core/agent_builder.rs | 74 +- crates/aether-core/src/core/agent_deps.rs | 33 - crates/aether-core/src/core/error.rs | 6 - crates/aether-core/src/core/prompt.rs | 70 +- crates/aether-core/src/core/tool_execution.rs | 122 ++- .../aether-core/src/docs/mcp_server_config.md | 64 ++ crates/aether-core/src/docs/tool_filter.md | 4 +- crates/aether-core/src/events/user_message.rs | 14 +- crates/aether-core/src/mcp/config.rs | 264 +++++++ crates/aether-core/src/mcp/deferred_tools.rs | 194 +++++ crates/aether-core/src/mcp/exposure.rs | 61 ++ crates/aether-core/src/mcp/gateway_service.rs | 187 ----- crates/aether-core/src/mcp/mcp_builder.rs | 702 +++++------------- crates/aether-core/src/mcp/mcp_catalog.rs | 202 +++++ crates/aether-core/src/mcp/mcp_handle.rs | 132 ---- crates/aether-core/src/mcp/mod.rs | 28 +- .../mcp/progressive_discovery_instructions.md | 16 + crates/aether-core/src/mcp/run_mcp_task.rs | 91 --- crates/aether-core/src/mcp/tool_bridge.rs | 83 ++- crates/aether-core/src/mcp/tool_filter.rs | 81 ++ crates/aether-core/src/testing/fake_mcp.rs | 30 - crates/aether-core/src/testing/mcp_test.rs | 283 ------- crates/aether-core/src/testing/mod.rs | 4 - crates/aether-core/src/testing/utils.rs | 79 +- crates/aether-core/tests/agent/agent_tests.rs | 41 +- .../aether-core/tests/agent/cancel_tests.rs | 95 +-- .../tests/agent/clear_context_tests.rs | 12 +- .../tests/agent/prompt_cache_tests.rs | 22 +- crates/aether-core/tests/agent/retry_tests.rs | 20 +- crates/aether-core/tests/agent/usage_tests.rs | 5 +- crates/aether-core/tests/doc_examples.rs | 2 +- crates/aether-core/tests/mcp.rs | 4 +- .../aether-core/tests/mcp/agent_sync_tests.rs | 184 ++--- .../tests/mcp/config_parser_tests.rs | 517 ++++++++----- .../tests/mcp/instructions_tests.rs | 81 +- crates/aether-core/tests/mcp/ipc_tests.rs | 217 ++++-- crates/aether-core/tests/mcp/mrtr_tests.rs | 133 +--- crates/aether-core/tests/mcp/oauth_tests.rs | 317 -------- .../tests/mcp/server_round_trip_tests.rs | 89 +-- crates/aether-core/tests/mcp/support.rs | 96 +++ crates/aether-core/tests/mcp/task_tests.rs | 403 ++-------- .../tests/mcp/tool_catalog_tests.rs | 212 ++++++ .../tests/mcp/tool_filter_tests.rs | 140 ++++ .../tests/mcp/trace_context_tests.rs | 26 +- .../docs/aether/settings/mcp-servers.mdx | 30 +- 52 files changed, 2638 insertions(+), 3206 deletions(-) create mode 100644 crates/aether-core/src/docs/mcp_server_config.md create mode 100644 crates/aether-core/src/mcp/config.rs create mode 100644 crates/aether-core/src/mcp/deferred_tools.rs create mode 100644 crates/aether-core/src/mcp/exposure.rs delete mode 100644 crates/aether-core/src/mcp/gateway_service.rs create mode 100644 crates/aether-core/src/mcp/mcp_catalog.rs delete mode 100644 crates/aether-core/src/mcp/mcp_handle.rs create mode 100644 crates/aether-core/src/mcp/progressive_discovery_instructions.md delete mode 100644 crates/aether-core/src/mcp/run_mcp_task.rs create mode 100644 crates/aether-core/src/mcp/tool_filter.rs delete mode 100644 crates/aether-core/src/testing/fake_mcp.rs delete mode 100644 crates/aether-core/src/testing/mcp_test.rs delete mode 100644 crates/aether-core/tests/mcp/oauth_tests.rs create mode 100644 crates/aether-core/tests/mcp/support.rs create mode 100644 crates/aether-core/tests/mcp/tool_catalog_tests.rs create mode 100644 crates/aether-core/tests/mcp/tool_filter_tests.rs diff --git a/Cargo.lock b/Cargo.lock index bcbb36f86..69f28224c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -134,6 +134,7 @@ dependencies = [ "thiserror 2.0.21", "tokio", "tokio-stream", + "tokio-util", "tracing", "tracing-subscriber", "uuid", 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-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/aether-core/src/docs/mcp_server_config.md b/crates/aether-core/src/docs/mcp_server_config.md new file mode 100644 index 000000000..015b170c1 --- /dev/null +++ b/crates/aether-core/src/docs/mcp_server_config.md @@ -0,0 +1,64 @@ +A single MCP server. The `type` field selects the transport; `stdio` is the +default and may be omitted. + +A local stdio server launched as a subprocess: + +```json +{ + "type": "stdio", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-github"], + "env": { "GITHUB_TOKEN": "ghp_..." } +} +``` + +A remote streamable-HTTP server using a pre-registered public OAuth client: + +```json +{ + "type": "http", + "url": "https://mcp.slack.com/mcp", + "oauth": { + "clientId": "1601185624273.8899143856786", + "callbackPort": 3118 + } +} +``` + +When a remote HTTP server returns an OAuth challenge, Aether uses its first-party +Client ID Metadata Document at +`https://aether-agent.io/oauth/client-metadata.json` and listens on +`127.0.0.1:3118`, advertising the exact redirect URI `http://localhost:3118/`. +If port 3118 is occupied, authentication fails until the port is available. + +The `oauth` object is optional. Set `clientMetadataUrl` for a custom CIMD, or +`clientId` for a pre-registered public client; a configured `clientId` takes +priority. `callbackPort` defaults to 3118 and must exactly match the client's +registered redirect URI. If the authorization server does not advertise CIMD, +Aether falls back to deprecated Dynamic Client Registration. An explicit +`Authorization` header bypasses OAuth entirely. + +A remote server using a bearer token: + +```json +{ + "type": "http", + "url": "https://mcp.example.com", + "headers": { "Authorization": "Bearer ..." } +} +``` + +Set `"deferTools": true` to omit every tool on this server from the model-visible +tool definitions and make it available through progressive `aether mcp` discovery. +For selective deferral, set `deferTools` to an object with `include` and `exclude` +lists. Entries match either an exact MCP-local tool name or a prefix ending in `*`: + +```json +{ + "type": "in-memory", + "deferTools": { + "include": ["*"], + "exclude": ["bash", "lsp_*"] + } +} +``` 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/aether-core/src/mcp/progressive_discovery_instructions.md b/crates/aether-core/src/mcp/progressive_discovery_instructions.md new file mode 100644 index 000000000..83d847211 --- /dev/null +++ b/crates/aether-core/src/mcp/progressive_discovery_instructions.md @@ -0,0 +1,16 @@ +You have additional MCP tools available through the `aether mcp` command via Bash. + +## Usage + +- `aether mcp --help` print connected MCP servers. +- `aether mcp --help` print a server's tools. +- `aether mcp --help` print the tool's schema. + +## Calling MCP tools + +Invocation accepts one JSON object only: + +- `aether mcp --json '{...}'` +- Pass a JSON object on stdin, for example `printf '%s' '{...}' | aether mcp `. + +The command prints one JSON value to stdout, so you can compose calls with `jq`, pipes, redirects, `&&`, scripts, etc. 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/aether-core/src/mcp/tool_filter.rs b/crates/aether-core/src/mcp/tool_filter.rs new file mode 100644 index 000000000..1ce23666c --- /dev/null +++ b/crates/aether-core/src/mcp/tool_filter.rs @@ -0,0 +1,81 @@ +use rmcp::model::Tool; +use utils::matches_name_pattern; + +#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] +#[serde(untagged)] +pub enum ToolMatcher { + Name(String), + Annotations(ToolAnnotationMatcher), +} + +impl ToolMatcher { + pub fn name(pattern: impl Into) -> Self { + Self::Name(pattern.into()) + } + + pub fn read_only() -> Self { + Self::Annotations(ToolAnnotationMatcher { read_only: Some(true), ..ToolAnnotationMatcher::default() }) + } + + pub fn annotations(matcher: ToolAnnotationMatcher) -> Self { + Self::Annotations(matcher) + } + + pub fn matches(&self, name: &str, tool: &Tool) -> bool { + match self { + Self::Name(pattern) => matches_name_pattern(pattern, name), + Self::Annotations(matcher) => matcher.matches(tool), + } + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ToolAnnotationMatcher { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub read_only: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub destructive: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub idempotent: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub open_world: Option, +} + +impl ToolAnnotationMatcher { + 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), + (self.destructive, annotations.destructive_hint), + (self.idempotent, annotations.idempotent_hint), + (self.open_world, annotations.open_world_hint), + ]; + if pairs.iter().all(|(field, _)| field.is_none()) { + return false; + } + pairs.iter().all(|(field, hint)| field.is_none_or(|value| *hint == Some(value))) + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ToolFilter { + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub allow: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub deny: Vec, +} + +impl ToolFilter { + pub fn is_empty(&self) -> bool { + self.allow.is_empty() && self.deny.is_empty() + } + + 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/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: From a3f8bcb9e50b52dea4dbd270b21fa5a5a3ee2908 Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:13 -0700 Subject: [PATCH 3/6] refactor(aether-project): centralize MCP source resolution Use aether-core MCP configuration and tool-filter types instead of depending directly on mcp-utils. Represent inline sources with McpConfig and centralize file and source resolution on McpFileSpec and McpSourceSpec. Reuse the same optional-file and variable handling when resolving agent catalogs and producing portable settings. --- crates/aether-project/Cargo.toml | 1 - crates/aether-project/src/aether_settings.rs | 90 ++++-------------- crates/aether-project/src/agent_catalog.rs | 93 +++++++------------ crates/aether-project/src/agent_config.rs | 2 +- crates/aether-project/src/error.rs | 7 ++ crates/aether-project/src/lib.rs | 4 +- ...ig_source_config.rs => mcp_source_spec.rs} | 61 ++++++++++-- crates/aether-project/src/testing.rs | 8 +- 8 files changed, 126 insertions(+), 140 deletions(-) rename crates/aether-project/src/{mcp_config_source_config.rs => mcp_source_spec.rs} (59%) 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(), } From 660af4ed51c66f41d87f7e0c56ece4b87b8d708a Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:13 -0700 Subject: [PATCH 4/6] refactor(mcp-servers): unify built-in server hosting Use McpServer for built-in server construction and share the Builtin and Host setup paths between agent-hosted and standalone servers. Adapt request context, background tasks, and subagent execution to the new MCP APIs. Pass deferred-tool IPC access through the agent's bash environment rather than owning discovery instructions here. Migrate server integration tests and documentation, and cover built-in construction and stdio transport through the new hosting interfaces. --- crates/aether-cli/Cargo.toml | 2 +- crates/mcp-servers/Cargo.toml | 21 +- crates/mcp-servers/README.md | 38 ++-- crates/mcp-servers/src/background_tasks.rs | 80 +++++++ crates/mcp-servers/src/bin/stdio.rs | 75 ++----- crates/mcp-servers/src/coding/mod.rs | 185 ++++++---------- crates/mcp-servers/src/docs/coding_mcp.md | 13 +- .../mcp-servers/src/docs/mcp_builder_ext.md | 8 +- crates/mcp-servers/src/docs/skills_mcp.md | 11 +- crates/mcp-servers/src/docs/subagents_mcp.md | 16 +- crates/mcp-servers/src/docs/tasks_mcp.md | 12 +- crates/mcp-servers/src/docs/tasks_module.md | 6 +- crates/mcp-servers/src/error.rs | 9 - crates/mcp-servers/src/lib.rs | 24 ++- .../src/progressive_discovery_instructions.md | 16 -- crates/mcp-servers/src/request_context.rs | 37 ++++ crates/mcp-servers/src/review/server.rs | 57 ++--- .../src/review/tools/review_artifact/mod.rs | 8 +- crates/mcp-servers/src/setup.rs | 168 +++++---------- crates/mcp-servers/src/skills/server.rs | 129 ++++------- crates/mcp-servers/src/subagents/server.rs | 114 +++------- .../src/subagents/tools/spawn_subagent/mod.rs | 113 +++++----- crates/mcp-servers/src/tasks/server.rs | 58 ++--- crates/mcp-servers/src/workspace_paths.rs | 11 + .../tests/integration/builtin_servers.rs | 42 ++++ .../tests/integration/coding_mcp_tools.rs | 42 ++-- .../tests/integration/coding_permissions.rs | 87 ++++++-- .../tests/integration/common/mod.rs | 204 +++++++----------- .../integration/lsp_check_errors_contract.rs | 21 +- .../tests/integration/lsp_diagnostics_e2e.rs | 26 ++- .../tests/integration/lsp_operations_e2e.rs | 11 +- .../integration/lsp_ts_diagnostics_e2e.rs | 14 +- .../integration/lsp_ts_operations_e2e.rs | 9 +- crates/mcp-servers/tests/integration/main.rs | 1 + .../tests/integration/plugins_mcp_agents.rs | 199 +++++++---------- .../tests/integration/plugins_mcp_skills.rs | 21 +- .../tests/integration/review_mcp.rs | 42 ++-- .../integration/skills_self_improvement.rs | 7 +- .../tests/integration/stdio_transport.rs | 101 +++------ .../tests/integration/test_bash.rs | 72 +++---- .../test_read_before_edit_safety.rs | 14 +- .../integration/test_read_file_rule_meta.rs | 4 +- .../tests/lsp_symbol_errors_e2e.rs | 8 +- .../tests/lsp_symbol_recovery_e2e.rs | 11 +- 44 files changed, 963 insertions(+), 1184 deletions(-) create mode 100644 crates/mcp-servers/src/background_tasks.rs delete mode 100644 crates/mcp-servers/src/error.rs delete mode 100644 crates/mcp-servers/src/progressive_discovery_instructions.md create mode 100644 crates/mcp-servers/src/request_context.rs create mode 100644 crates/mcp-servers/tests/integration/builtin_servers.rs diff --git a/crates/aether-cli/Cargo.toml b/crates/aether-cli/Cargo.toml index 785201fb3..e7a2af84f 100644 --- a/crates/aether-cli/Cargo.toml +++ b/crates/aether-cli/Cargo.toml @@ -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" } 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/progressive_discovery_instructions.md b/crates/mcp-servers/src/progressive_discovery_instructions.md deleted file mode 100644 index 83d847211..000000000 --- a/crates/mcp-servers/src/progressive_discovery_instructions.md +++ /dev/null @@ -1,16 +0,0 @@ -You have additional MCP tools available through the `aether mcp` command via Bash. - -## Usage - -- `aether mcp --help` print connected MCP servers. -- `aether mcp --help` print a server's tools. -- `aether mcp --help` print the tool's schema. - -## Calling MCP tools - -Invocation accepts one JSON object only: - -- `aether mcp --json '{...}'` -- Pass a JSON object on stdin, for example `printf '%s' '{...}' | aether mcp `. - -The command prints one JSON value to stdout, so you can compose calls with `jq`, pipes, redirects, `&&`, scripts, etc. 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()); } From 2a59ed9d04390cb1814349eadd986471c497d999 Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:14 -0700 Subject: [PATCH 5/6] test(aether-sessions): compare non-system conversation identities Exclude system messages when comparing live and replayed message identities so session assertions focus on durable conversation boundaries rather than runtime-managed prompt content. Share the identity comparison helper and compare the continuation context against the corresponding restored conversation prefix. --- crates/aether-sessions/tests/message_identity.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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() +} From 98eed0d8febdea6d34e90664e15368c9172f15f4 Mon Sep 17 00:00:00 2001 From: Josh Carver Date: Mon, 12 Oct 2026 06:53:14 -0700 Subject: [PATCH 6/6] refactor(aether-cli): migrate runtimes and ACP to MCP catalogs Use McpCatalog throughout runtime construction, headless execution, prompt previews, MCP commands, and ACP sessions instead of the legacy MCP runtime and event-pump APIs. Resolve slash commands from server-associated prompts and keep ACP command formatting in the session layer. Exercise prompt expansion and command updates through ACP integration tests. Preserve typed selection, configuration, startup, and agent errors. Update generated settings and documentation for the new MCP config shape, and refresh the workspace lockfile. --- .aether/settings.json | 3 +- Cargo.lock | 15 +- crates/aether-cli/Cargo.toml | 3 +- crates/aether-cli/README.md | 2 +- crates/aether-cli/src/acp/fake_prompt_mcp.rs | 83 ----- crates/aether-cli/src/acp/mod.rs | 3 - .../aether-cli/src/acp/protocol/commands.rs | 19 +- crates/aether-cli/src/acp/protocol/mcp.rs | 99 ++--- crates/aether-cli/src/acp/protocol/mod.rs | 2 - crates/aether-cli/src/acp/session/actor.rs | 337 ++++-------------- crates/aether-cli/src/acp/session/error.rs | 11 +- crates/aether-cli/src/acp/session/factory.rs | 27 +- crates/aether-cli/src/acp/session/mod.rs | 1 - crates/aether-cli/src/acp/session/model.rs | 4 +- crates/aether-cli/src/acp/session/runtime.rs | 121 +++---- .../src/acp/session/slash_commands.rs | 39 -- crates/aether-cli/src/acp/state.rs | 48 ++- crates/aether-cli/src/acp/testing.rs | 298 ++++++++-------- crates/aether-cli/src/error.rs | 17 +- crates/aether-cli/src/headless/mod.rs | 35 +- crates/aether-cli/src/headless/run.rs | 59 ++- crates/aether-cli/src/init/build_settings.rs | 7 +- crates/aether-cli/src/lib.rs | 4 +- crates/aether-cli/src/mcp_command.rs | 144 +++----- crates/aether-cli/src/mcp_config_args.rs | 13 +- crates/aether-cli/src/resolve.rs | 9 +- crates/aether-cli/src/runtime.rs | 215 +++++------ crates/aether-cli/src/show_prompt/run.rs | 30 +- crates/aether-cli/src/slash_commands.rs | 137 ++++--- .../tests/integration/acp_agent_switching.rs | 6 +- .../tests/integration/acp_mcp_sign_in.rs | 46 +++ .../tests/integration/acp_prompt_lifecycle.rs | 45 ++- .../tests/integration/acp_remote.rs | 45 ++- .../integration/acp_session_lifecycle.rs | 15 + .../tests/integration/acp_slash_commands.rs | 108 ++++++ crates/aether-cli/tests/integration/main.rs | 3 +- .../tests/integration/mcp_command.rs | 145 ++++---- .../aether-cli/tests/integration/run_init.rs | 49 +-- .../tests/integration/slash_commands.rs | 49 --- 39 files changed, 972 insertions(+), 1324 deletions(-) delete mode 100644 crates/aether-cli/src/acp/fake_prompt_mcp.rs delete mode 100644 crates/aether-cli/src/acp/session/slash_commands.rs create mode 100644 crates/aether-cli/tests/integration/acp_mcp_sign_in.rs create mode 100644 crates/aether-cli/tests/integration/acp_slash_commands.rs delete mode 100644 crates/aether-cli/tests/integration/slash_commands.rs 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/Cargo.lock b/Cargo.lock index 69f28224c..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", @@ -333,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", ] @@ -361,7 +356,6 @@ dependencies = [ "aether-agent-core", "aether-doctest", "aether-llm", - "aether-mcp-utils", "aether-project", "aether-utils", "globset", @@ -6619,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", @@ -6634,6 +6631,7 @@ dependencies = [ "tokio", "tokio-stream", "tokio-util", + "tower-service", "tracing", "url", "uuid", @@ -7929,6 +7927,7 @@ dependencies = [ "futures-core", "pin-project-lite", "tokio", + "tokio-util", ] [[package]] diff --git a/crates/aether-cli/Cargo.toml b/crates/aether-cli/Cargo.toml index e7a2af84f..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"] } @@ -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"), - } -}