| 1 | //! Streamable HTTP against an in-process rmcp server, and the Ed executor on top of it. |
| 2 | |
| 3 | use std::sync::Arc; |
| 4 | use std::time::Duration; |
| 5 | |
| 6 | use axum::extract::Request; |
| 7 | use axum::http::StatusCode; |
| 8 | use axum::middleware::{self, Next}; |
| 9 | use axum::response::Response; |
| 10 | use ed_agent::{ToolExecutor, ToolRisk}; |
| 11 | use ed_mcp::{ |
| 12 | ClientInfo, Connection, ElicitRequestParams, ElicitResult, Elicitation, ElicitationAction, |
| 13 | McpExecutor, McpToolset, ServerSpec, ToolPolicy, ToolsetConfig, |
| 14 | }; |
| 15 | use rmcp::RoleServer; |
| 16 | use rmcp::ServerHandler; |
| 17 | use rmcp::model::{ |
| 18 | CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ListToolsResult, |
| 19 | ServerCapabilities, ServerConfig, Tool, ToolAnnotations, |
| 20 | }; |
| 21 | use rmcp::service::RequestContext; |
| 22 | use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; |
| 23 | use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; |
| 24 | use serde_json::json; |
| 25 | use tokio_util::sync::CancellationToken; |
| 26 | |
| 27 | const TOKEN: &str = "s3cret"; |
| 28 | |
| 29 | #[derive(Clone)] |
| 30 | struct Mailbox; |
| 31 | |
| 32 | impl ServerHandler for Mailbox { |
| 33 | fn get_info(&self) -> ServerConfig { |
| 34 | ServerConfig::new(ServerCapabilities::builder().enable_tools().build()) |
| 35 | } |
| 36 | |
| 37 | async fn list_tools( |
| 38 | &self, |
| 39 | _request: Option<rmcp::model::PaginatedRequestParams>, |
| 40 | _context: RequestContext<RoleServer>, |
| 41 | ) -> Result<ListToolsResult, rmcp::ErrorData> { |
| 42 | let schema = Arc::new(json!({"type": "object"}).as_object().unwrap().clone()); |
| 43 | Ok(ListToolsResult { |
| 44 | tools: vec![ |
| 45 | // Declares itself read-only. |
| 46 | Tool::new("whoami", "Who is signed in.", schema.clone()) |
| 47 | .annotate(ToolAnnotations::new().read_only(true)), |
| 48 | // Says nothing, like a server that predates annotations. |
| 49 | Tool::new("list_messages", "Recent messages.", schema.clone()), |
| 50 | Tool::new("send_message", "Send mail.", schema.clone()), |
| 51 | // Elicits structured input from the user mid-call. |
| 52 | Tool::new("collect_feedback", "Ask the user for feedback.", schema), |
| 53 | ], |
| 54 | ..Default::default() |
| 55 | }) |
| 56 | } |
| 57 | |
| 58 | async fn call_tool( |
| 59 | &self, |
| 60 | request: CallToolRequestParams, |
| 61 | context: RequestContext<RoleServer>, |
| 62 | ) -> Result<CallToolResponse, rmcp::ErrorData> { |
| 63 | // `collect_feedback` asks the client (via `elicitation/create`) for input and reports |
| 64 | // what came back, so a test can prove the request reached the client's hook and the |
| 65 | // answer flowed back. |
| 66 | if request.name == "collect_feedback" { |
| 67 | let params: ElicitRequestParams = serde_json::from_value(json!({ |
| 68 | "mode": "form", |
| 69 | "message": "How was it?", |
| 70 | "requestedSchema": { |
| 71 | "type": "object", |
| 72 | "properties": {"rating": {"type": "string"}}, |
| 73 | "required": ["rating"] |
| 74 | } |
| 75 | })) |
| 76 | .unwrap(); |
| 77 | let result = context |
| 78 | .peer |
| 79 | .create_elicitation(params) |
| 80 | .await |
| 81 | .map_err(|e| rmcp::ErrorData::internal_error(e.to_string(), None))?; |
| 82 | let text = match result.action { |
| 83 | ElicitationAction::Accept => { |
| 84 | let rating = result |
| 85 | .content |
| 86 | .as_ref() |
| 87 | .and_then(|c: &serde_json::Value| c.get("rating")) |
| 88 | .and_then(|v: &serde_json::Value| v.as_str()) |
| 89 | .unwrap_or(""); |
| 90 | format!("accepted: {rating}") |
| 91 | } |
| 92 | ElicitationAction::Decline => "declined".to_string(), |
| 93 | ElicitationAction::Cancel => "cancelled".to_string(), |
| 94 | _ => "unknown".to_string(), |
| 95 | }; |
| 96 | return Ok(CallToolResult::success(vec![ContentBlock::text(text)]).into()); |
| 97 | } |
| 98 | Ok( |
| 99 | CallToolResult::success(vec![ContentBlock::text(format!("ran {}", request.name))]) |
| 100 | .into(), |
| 101 | ) |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | async fn require_bearer(request: Request, next: Next) -> Result<Response, StatusCode> { |
| 106 | let ok = request |
| 107 | .headers() |
| 108 | .get("authorization") |
| 109 | .and_then(|v| v.to_str().ok()) |
| 110 | == Some(&format!("Bearer {TOKEN}")); |
| 111 | if ok { |
| 112 | Ok(next.run(request).await) |
| 113 | } else { |
| 114 | Err(StatusCode::UNAUTHORIZED) |
| 115 | } |
| 116 | } |
| 117 | |
| 118 | /// Serve [`Mailbox`] behind bearer auth; returns its URL. |
| 119 | async fn serve() -> String { |
| 120 | let config = StreamableHttpServerConfig::default() |
| 121 | .with_legacy_session_mode(true) |
| 122 | .with_sse_keep_alive(None) |
| 123 | .with_cancellation_token(CancellationToken::new()); |
| 124 | let service: StreamableHttpService<Mailbox, LocalSessionManager> = |
| 125 | StreamableHttpService::new(|| Ok(Mailbox), Default::default(), config); |
| 126 | let router = axum::Router::new() |
| 127 | .nest_service("/mcp", service) |
| 128 | .layer(middleware::from_fn(require_bearer)); |
| 129 | let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); |
| 130 | let addr = listener.local_addr().unwrap(); |
| 131 | tokio::spawn(async move { |
| 132 | let _ = axum::serve(listener, router).await; |
| 133 | }); |
| 134 | format!("http://{addr}/mcp") |
| 135 | } |
| 136 | |
| 137 | fn client() -> ClientInfo { |
| 138 | ClientInfo::new("ed-mcp-test", "0") |
| 139 | } |
| 140 | |
| 141 | #[tokio::test] |
| 142 | async fn connects_lists_and_calls_with_a_bearer_token() { |
| 143 | let url = serve().await; |
| 144 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 145 | |
| 146 | let set = McpToolset::connect_specs(ToolsetConfig::new(client()), &[spec], &[]).await; |
| 147 | |
| 148 | let names: Vec<_> = set.tools().iter().map(|t| t.name.clone()).collect(); |
| 149 | assert_eq!( |
| 150 | names, |
| 151 | [ |
| 152 | "mcp__mail__collect_feedback", |
| 153 | "mcp__mail__list_messages", |
| 154 | "mcp__mail__send_message", |
| 155 | "mcp__mail__whoami" |
| 156 | ] |
| 157 | ); |
| 158 | let out = set |
| 159 | .call( |
| 160 | "mcp__mail__whoami", |
| 161 | json!({}), |
| 162 | &CancellationToken::new(), |
| 163 | &|_| {}, |
| 164 | ) |
| 165 | .await |
| 166 | .unwrap(); |
| 167 | assert_eq!(out.text, "ran whoami"); |
| 168 | assert!(!out.failed); |
| 169 | set.shutdown().await; |
| 170 | } |
| 171 | |
| 172 | #[tokio::test] |
| 173 | async fn a_connection_lists_and_calls_tools_by_their_own_names() { |
| 174 | let url = serve().await; |
| 175 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 176 | let conn = Connection::connect(&spec, &client(), Duration::from_secs(5)) |
| 177 | .await |
| 178 | .unwrap(); |
| 179 | |
| 180 | let mut names: Vec<_> = conn.tools().iter().map(|t| t.name.to_string()).collect(); |
| 181 | names.sort(); |
| 182 | assert_eq!( |
| 183 | names, |
| 184 | [ |
| 185 | "collect_feedback", |
| 186 | "list_messages", |
| 187 | "send_message", |
| 188 | "whoami" |
| 189 | ] |
| 190 | ); |
| 191 | |
| 192 | let result = conn |
| 193 | .call_tool("whoami", json!({}), Duration::from_secs(5)) |
| 194 | .await |
| 195 | .unwrap(); |
| 196 | assert_eq!(result["content"][0]["text"], "ran whoami"); |
| 197 | let err = conn |
| 198 | .call_tool("whoami", json!([1]), Duration::from_secs(5)) |
| 199 | .await |
| 200 | .unwrap_err(); |
| 201 | assert!(err.to_string().contains("JSON object"), "{err:#}"); |
| 202 | } |
| 203 | |
| 204 | #[tokio::test] |
| 205 | async fn a_wrong_bearer_token_is_rejected() { |
| 206 | let url = serve().await; |
| 207 | let spec = ServerSpec::http("mail", url).with_bearer("nope"); |
| 208 | |
| 209 | let err = Connection::connect(&spec, &client(), Duration::from_secs(5)) |
| 210 | .await |
| 211 | .err() |
| 212 | .expect("the server answers 401"); |
| 213 | assert!(err.to_string().contains("handshake failed"), "{err:#}"); |
| 214 | |
| 215 | // A rejected server is left out, not fatal. |
| 216 | let set = McpToolset::connect_specs(ToolsetConfig::new(client()), &[spec], &[]).await; |
| 217 | assert!(set.tools().is_empty()); |
| 218 | } |
| 219 | |
| 220 | #[tokio::test] |
| 221 | async fn a_missing_bearer_token_is_rejected() { |
| 222 | let url = serve().await; |
| 223 | let spec = ServerSpec::http("mail", url); |
| 224 | assert!( |
| 225 | Connection::connect(&spec, &client(), Duration::from_secs(5)) |
| 226 | .await |
| 227 | .is_err() |
| 228 | ); |
| 229 | } |
| 230 | |
| 231 | async fn executor(policy: ToolPolicy) -> McpExecutor { |
| 232 | let url = serve().await; |
| 233 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 234 | let set = McpToolset::connect_specs(ToolsetConfig::new(client()), &[spec], &[]).await; |
| 235 | let exec = McpExecutor::new(policy); |
| 236 | exec.set_toolset(Some(set)).await; |
| 237 | exec |
| 238 | } |
| 239 | |
| 240 | fn risk(defs: &[ed_agent::AgentToolDefinition], name: &str) -> ToolRisk { |
| 241 | defs.iter().find(|d| d.name == name).unwrap().risk |
| 242 | } |
| 243 | |
| 244 | #[tokio::test] |
| 245 | async fn risk_follows_the_read_only_hint() { |
| 246 | let defs = executor(ToolPolicy::new()).await.definitions().await; |
| 247 | |
| 248 | assert_eq!(risk(&defs, "mcp__mail__whoami"), ToolRisk::ReadOnly); |
| 249 | // No hint means it may change things. |
| 250 | assert_eq!(risk(&defs, "mcp__mail__list_messages"), ToolRisk::Mutating); |
| 251 | assert_eq!(risk(&defs, "mcp__mail__send_message"), ToolRisk::Mutating); |
| 252 | } |
| 253 | |
| 254 | #[tokio::test] |
| 255 | async fn the_host_can_override_risk_and_drop_the_prefix() { |
| 256 | let policy = ToolPolicy::new() |
| 257 | .read_only(&["list_messages", "whoami"]) |
| 258 | .unprefixed(); |
| 259 | let exec = executor(policy).await; |
| 260 | let defs = exec.definitions().await; |
| 261 | |
| 262 | assert_eq!(risk(&defs, "list_messages"), ToolRisk::ReadOnly); |
| 263 | assert_eq!(risk(&defs, "whoami"), ToolRisk::ReadOnly); |
| 264 | assert_eq!(risk(&defs, "send_message"), ToolRisk::Mutating); |
| 265 | assert_eq!( |
| 266 | defs.iter() |
| 267 | .find(|d| d.name == "send_message") |
| 268 | .unwrap() |
| 269 | .description, |
| 270 | "Send mail." |
| 271 | ); |
| 272 | |
| 273 | // The executor answers to the names the model sees. |
| 274 | assert_eq!( |
| 275 | exec.execute("send_message", "{}").await.unwrap(), |
| 276 | "ran send_message" |
| 277 | ); |
| 278 | assert!(exec.execute("mcp__mail__send_message", "{}").await.is_err()); |
| 279 | } |
| 280 | |
| 281 | /// Records the requests it is handed and answers each with a fixed action. |
| 282 | struct Recorder { |
| 283 | seen: Arc<std::sync::Mutex<Vec<(String, ElicitRequestParams)>>>, |
| 284 | answer: ElicitResult, |
| 285 | } |
| 286 | |
| 287 | impl Elicitation for Recorder { |
| 288 | fn elicit( |
| 289 | &self, |
| 290 | server: &str, |
| 291 | request: ElicitRequestParams, |
| 292 | ) -> futures::future::BoxFuture<'_, anyhow::Result<ElicitResult>> { |
| 293 | self.seen |
| 294 | .lock() |
| 295 | .unwrap() |
| 296 | .push((server.to_string(), request)); |
| 297 | let answer = self.answer.clone(); |
| 298 | Box::pin(async move { Ok(answer) }) |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | #[tokio::test] |
| 303 | async fn forwards_server_elicitation_to_the_hook() { |
| 304 | let url = serve().await; |
| 305 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 306 | |
| 307 | let seen = Arc::new(std::sync::Mutex::new(Vec::new())); |
| 308 | let hook = Arc::new(Recorder { |
| 309 | seen: seen.clone(), |
| 310 | answer: ElicitResult::new(ElicitationAction::Accept) |
| 311 | .with_content(json!({"rating": "great"})), |
| 312 | }); |
| 313 | let config = ToolsetConfig::new(client()).with_elicitation(hook); |
| 314 | let set = McpToolset::connect_specs(config, &[spec], &[]).await; |
| 315 | |
| 316 | let out = set |
| 317 | .call( |
| 318 | "mcp__mail__collect_feedback", |
| 319 | json!({}), |
| 320 | &CancellationToken::new(), |
| 321 | &|_| {}, |
| 322 | ) |
| 323 | .await |
| 324 | .unwrap(); |
| 325 | |
| 326 | // The server's elicitation reached our hook, and the Accept + content flowed back to it. |
| 327 | assert!(!out.failed, "{}", out.text); |
| 328 | assert_eq!(out.text, "accepted: great"); |
| 329 | |
| 330 | // Scoped so the guard is dropped before the `await` below. |
| 331 | { |
| 332 | let seen = seen.lock().unwrap(); |
| 333 | assert_eq!(seen.len(), 1, "the hook was called exactly once"); |
| 334 | assert_eq!(seen[0].0, "mail", "the hook knows which server asked"); |
| 335 | match &seen[0].1 { |
| 336 | ElicitRequestParams::FormElicitationParams { message, .. } => { |
| 337 | assert_eq!(message, "How was it?"); |
| 338 | } |
| 339 | other => panic!("expected a form elicitation, got {other:?}"), |
| 340 | } |
| 341 | } |
| 342 | |
| 343 | set.shutdown().await; |
| 344 | } |
| 345 | |
| 346 | #[tokio::test] |
| 347 | async fn a_hook_can_cancel_an_elicitation() { |
| 348 | let url = serve().await; |
| 349 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 350 | |
| 351 | // A hook that cancels round-trips distinctly from one that declines: `create_elicitation` |
| 352 | // forwards the hook's `ElicitResult` verbatim, so the server sees Cancel, not Decline. |
| 353 | let seen = Arc::new(std::sync::Mutex::new(Vec::new())); |
| 354 | let hook = Arc::new(Recorder { |
| 355 | seen: seen.clone(), |
| 356 | answer: ElicitResult::new(ElicitationAction::Cancel), |
| 357 | }); |
| 358 | let config = ToolsetConfig::new(client()).with_elicitation(hook); |
| 359 | let set = McpToolset::connect_specs(config, &[spec], &[]).await; |
| 360 | |
| 361 | let out = set |
| 362 | .call( |
| 363 | "mcp__mail__collect_feedback", |
| 364 | json!({}), |
| 365 | &CancellationToken::new(), |
| 366 | &|_| {}, |
| 367 | ) |
| 368 | .await |
| 369 | .unwrap(); |
| 370 | |
| 371 | // The hook ran and the server mapped Cancel to "cancelled", distinct from "declined". |
| 372 | assert!(!out.failed, "{}", out.text); |
| 373 | assert_eq!(out.text, "cancelled"); |
| 374 | assert_eq!(seen.lock().unwrap().len(), 1, "the hook was called once"); |
| 375 | |
| 376 | set.shutdown().await; |
| 377 | } |
| 378 | |
| 379 | #[tokio::test] |
| 380 | async fn without_a_hook_server_elicitation_is_declined() { |
| 381 | let url = serve().await; |
| 382 | let spec = ServerSpec::http("mail", url).with_bearer(TOKEN); |
| 383 | |
| 384 | // No `with_elicitation`: the default declines, and the server sees a Decline. |
| 385 | let set = McpToolset::connect_specs(ToolsetConfig::new(client()), &[spec], &[]).await; |
| 386 | let out = set |
| 387 | .call( |
| 388 | "mcp__mail__collect_feedback", |
| 389 | json!({}), |
| 390 | &CancellationToken::new(), |
| 391 | &|_| {}, |
| 392 | ) |
| 393 | .await |
| 394 | .unwrap(); |
| 395 | assert_eq!(out.text, "declined"); |
| 396 | set.shutdown().await; |
| 397 | } |
| 398 | |
| 399 | #[tokio::test] |
| 400 | async fn disconnecting_removes_the_tools() { |
| 401 | let exec = executor(ToolPolicy::new()).await; |
| 402 | assert!(!exec.definitions().await.is_empty()); |
| 403 | |
| 404 | exec.set_toolset(None).await; |
| 405 | |
| 406 | assert!(exec.definitions().await.is_empty()); |
| 407 | assert!(exec.execute("mcp__mail__whoami", "{}").await.is_err()); |
| 408 | } |