aether_core/testing/
mcp_test.rs

1use crate::events::{TaskOutcomeState, TraceContext, task_created_result};
2use crate::mcp::tool_bridge::{convert_tool_result, map_task_result_to_outcome};
3use crate::mcp::{McpHandle, McpRuntime, ServerFactory, ToolCallStream, mcp};
4use futures::{FutureExt, StreamExt};
5use mcp_utils::client::{
6    CallToolOptions, CancellationToken, InMemoryServerSpec, McpConnectionDetails, McpServer, McpTransport,
7    ToolCallEvent, ToolExposure, ToolFilter,
8};
9use mcp_utils::testing::ElicitationScript;
10use rmcp::model::{CreateTaskResult, ElicitResult, ProgressNotificationParam};
11use rmcp::{RoleServer, ServerHandler, service::DynService};
12use serde_json::Value;
13use std::collections::{HashMap, VecDeque};
14use std::path::Path;
15use std::sync::Mutex;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::time::Duration;
18use tokio::sync::watch;
19
20pub use mcp_utils::testing::CapturedElicitation;
21
22const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(10);
23
24#[derive(Default)]
25pub struct McpTestBuilder {
26    servers: Vec<McpServer>,
27    factories: Vec<(String, ServerFactory)>,
28    elicitation_responses: Vec<ElicitResult>,
29    trace_context: Option<TraceContext>,
30    tool_timeout: Duration,
31    tool_filter: ToolFilter,
32}
33
34fn task_outcome(outcome: crate::events::TaskOutcome) -> TaskOutcome {
35    let (status, body) = match outcome.state {
36        TaskOutcomeState::Completed { result, .. } => ("completed", result.result),
37        TaskOutcomeState::Failed { error } => ("failed", error.error),
38        TaskOutcomeState::Cancelled => {
39            ("cancelled", "The background task was cancelled and will not produce a result.".into())
40        }
41    };
42    TaskOutcome { task_id: outcome.task_id, status: status.into(), body }
43}
44
45pub struct McpTest {
46    mcp: McpHandle,
47    runtime: McpRuntime,
48    snapshot: McpConnectionDetails,
49    elicitations: ElicitationScript,
50    deferred_tools: tokio::sync::Mutex<VecDeque<DeferredTool>>,
51    cancel_tokens: Mutex<HashMap<String, CancellationToken>>,
52    trace_context: Option<TraceContext>,
53    tool_timeout: Duration,
54    next_call_id: AtomicU64,
55}
56
57pub struct TaskOutcome {
58    pub task_id: String,
59    pub status: String,
60    pub body: String,
61}
62
63pub struct ToolCallOutcome {
64    pub result: Result<llm::ToolCallResult, llm::ToolCallError>,
65    pub progress: Vec<ProgressNotificationParam>,
66    pub deferred_task: Option<CreateTaskResult>,
67}
68
69struct DeferredTool {
70    request: llm::ToolCallRequest,
71    events: ToolCallStream,
72}
73
74impl McpTestBuilder {
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    pub fn server<S>(self, name: impl Into<String>, server: S) -> Self
80    where
81        S: ServerHandler + Clone + Send + Sync + 'static,
82    {
83        self.server_with_exposure(name, server, ToolExposure::ModelVisible)
84    }
85
86    pub fn deferred_server<T>(self, name: impl Into<String>, server: T) -> Self
87    where
88        T: ServerHandler + Clone + Send + Sync + 'static,
89    {
90        self.server_with_exposure(name, server, ToolExposure::deferred_all())
91    }
92
93    pub fn server_with_exposure<T>(mut self, name: impl Into<String>, server: T, exposure: ToolExposure) -> Self
94    where
95        T: ServerHandler + Clone + Send + Sync + 'static,
96    {
97        let name = name.into();
98        let factory_name = format!("test-{}", self.factories.len());
99        let factory_server = server;
100        let factory: ServerFactory = Box::new(move |_spec, _services| {
101            let server = factory_server.clone();
102            async move { Box::new(server) as Box<dyn DynService<RoleServer>> }.boxed()
103        });
104        self.factories.push((factory_name.clone(), factory));
105        self.servers.push(McpServer::new(
106            name,
107            McpTransport::InMemory {
108                spec: InMemoryServerSpec { factory: factory_name, args: Vec::new(), input: None },
109            },
110            exposure,
111        ));
112        self
113    }
114
115    pub fn elicitation_response(mut self, response: ElicitResult) -> Self {
116        self.elicitation_responses.push(response);
117        self
118    }
119
120    pub fn trace_context(mut self, trace_context: TraceContext) -> Self {
121        self.trace_context = Some(trace_context);
122        self
123    }
124
125    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
126        self.tool_timeout = timeout;
127        self
128    }
129
130    /// Filter which tools the servers expose and the gateway serves.
131    pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
132        self.tool_filter = filter;
133        self
134    }
135
136    pub async fn build(self) -> McpTest {
137        let builder = mcp("/workspace").with_servers(self.servers).with_tool_filter(self.tool_filter);
138        let builder = self
139            .factories
140            .into_iter()
141            .fold(builder, |builder, (name, factory)| builder.register_in_memory_server(name, factory));
142        let mut spawn = builder.spawn().await.expect("MCP test manager spawns");
143        let snapshot = spawn.block_until_ready().await.expect("MCP test manager becomes ready");
144        let (runtime, event_rx) = spawn.split();
145
146        McpTest {
147            mcp: runtime.handle().clone(),
148            runtime,
149            snapshot,
150            elicitations: ElicitationScript::spawn(event_rx, self.elicitation_responses),
151            deferred_tools: tokio::sync::Mutex::new(VecDeque::new()),
152            cancel_tokens: Mutex::new(HashMap::new()),
153            trace_context: self.trace_context,
154            tool_timeout: if self.tool_timeout.is_zero() { DEFAULT_TOOL_TIMEOUT } else { self.tool_timeout },
155            next_call_id: AtomicU64::new(1),
156        }
157    }
158}
159
160impl McpTest {
161    pub async fn call(&self, server: &str, tool: &str, arguments: Value) -> ToolCallOutcome {
162        let id = self.next_call_id.fetch_add(1, Ordering::Relaxed);
163        let request = llm::ToolCallRequest {
164            id: format!("mcp-test-{id}"),
165            name: format!("{server}__{tool}"),
166            arguments: arguments.to_string(),
167        };
168        let request_for_outcome = request.clone();
169        let cancel = CancellationToken::new();
170        self.cancel_tokens.lock().expect("cancel token lock").insert(request.id.clone(), cancel.clone());
171        let options = CallToolOptions {
172            timeout: self.tool_timeout,
173            meta: self.trace_context.as_ref().map(TraceContext::to_meta),
174            cancel,
175        };
176        let mut events = self.mcp.call_model_visible(request.name, &request.arguments, options);
177
178        let mut progress = Vec::new();
179        while let Some(event) = events.next().await {
180            match event {
181                ToolCallEvent::Progress(event) => progress.push(event),
182                ToolCallEvent::TaskCreated(task) => {
183                    self.deferred_tools
184                        .lock()
185                        .await
186                        .push_back(DeferredTool { request: request_for_outcome.clone(), events });
187                    return ToolCallOutcome {
188                        result: Ok(task_created_result(&request_for_outcome, &task.task.task_id)),
189                        progress,
190                        deferred_task: Some(task),
191                    };
192                }
193                ToolCallEvent::Complete(outcome) => {
194                    let result = convert_tool_result(&request_for_outcome, outcome).map(|(result, _)| result);
195                    return ToolCallOutcome { result, progress, deferred_task: None };
196                }
197                ToolCallEvent::TaskStatus(_) | ToolCallEvent::TaskComplete { .. } | ToolCallEvent::Cancelled { .. } => {
198                    panic!("MCP task lifecycle event arrived before deferral")
199                }
200            }
201        }
202        panic!("MCP test tool event stream ended before completion");
203    }
204
205    pub fn cancel_tool(&self, tool_id: &str) {
206        let tokens = self.cancel_tokens.lock().expect("cancel token lock");
207        tokens.get(tool_id).expect("cancel_tool targets a tool started with call()").cancel();
208    }
209
210    pub async fn next_tool_event(&self) -> Option<ToolCallEvent> {
211        self.next_deferred_event().await.map(|(_, event)| event)
212    }
213
214    pub async fn next_task_outcome(&self) -> Option<TaskOutcome> {
215        while let Some((request, event)) = self.next_deferred_event().await {
216            match event {
217                ToolCallEvent::TaskComplete { task, result } => {
218                    return Some(task_outcome(map_task_result_to_outcome(request, task, result)));
219                }
220                ToolCallEvent::Cancelled { task_id } => {
221                    return Some(task_outcome(crate::events::TaskOutcome {
222                        request,
223                        task_id: task_id.unwrap_or_else(|| "pending".to_string()),
224                        state: TaskOutcomeState::Cancelled,
225                    }));
226                }
227                ToolCallEvent::Progress(_)
228                | ToolCallEvent::TaskCreated(_)
229                | ToolCallEvent::TaskStatus(_)
230                | ToolCallEvent::Complete(_) => {}
231            }
232        }
233        None
234    }
235
236    async fn next_deferred_event(&self) -> Option<(llm::ToolCallRequest, ToolCallEvent)> {
237        let mut deferred = self.deferred_tools.lock().await;
238        loop {
239            let front = deferred.front_mut()?;
240            match front.events.next().await {
241                Some(event) => return Some((front.request.clone(), event)),
242                None => {
243                    deferred.pop_front();
244                }
245            }
246        }
247    }
248
249    pub fn snapshot(&self) -> &McpConnectionDetails {
250        &self.snapshot
251    }
252
253    /// The Unix socket path of the deferred-tool gateway, if one was started.
254    pub fn gateway_endpoint(&self) -> Option<&Path> {
255        self.runtime.gateway_endpoint()
256    }
257
258    pub fn subscribe(&self) -> watch::Receiver<McpConnectionDetails> {
259        self.mcp.subscribe()
260    }
261
262    pub fn elicitations(&self) -> Vec<CapturedElicitation> {
263        self.elicitations.captured()
264    }
265}
aether_core/testing/mcp_test.rs

aether_core/testing/
mcp_test.rs
