agent_sdk_toolkit/testing/protocol/
mcp.rs

1//! Scripted MCP protocol harnesses for SDK consumers. Use these fakes to test
2//! tool/resource/prompt lifecycle and host filtering without live MCP servers.
3//! Harness methods mutate in-memory endpoints and scripted response state only.
4//!
5use std::collections::{BTreeMap, BTreeSet};
6
7use agent_sdk_core::{AgentError, AgentErrorKind, RetryClassification};
8use serde_json::{Value, json};
9
10use crate::protocol::{
11    JsonRpcFrame, JsonRpcId, JsonRpcLineEndpoint, JsonRpcRequest, JsonRpcResponse, expect_response,
12    protocol_violation,
13};
14
15enum ReceiveNext {
16    Empty,
17    Handled,
18    Frame(JsonRpcFrame),
19}
20
21#[derive(Clone, Debug, Default)]
22/// In-memory scripted mcp server fixture for SDK conformance tests.
23/// Use it to script deterministic behavior in memory; any transcript or endpoint mutation is documented on the method that performs it.
24pub struct ScriptedMcpServer {
25    tools: BTreeMap<String, Value>,
26    resources: BTreeMap<String, String>,
27    prompts: BTreeMap<String, Value>,
28    initialized: bool,
29    next_client_request: i64,
30}
31
32impl ScriptedMcpServer {
33    /// Creates a new testing::protocol::mcp value with explicit
34    /// caller-provided inputs. This constructor is data-only and
35    /// performs no I/O or external side effects.
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Returns an updated testing::protocol::mcp value with tool applied.
41    /// This is data construction only and does not execute the configured
42    /// behavior.
43    pub fn tool(mut self, name: impl Into<String>, result: Value) -> Self {
44        self.tools.insert(name.into(), result);
45        self
46    }
47
48    /// Returns an updated value with resource configured.
49    /// This configures or reads the scripted protocol mock in memory.
50    pub fn resource(mut self, uri: impl Into<String>, text: impl Into<String>) -> Self {
51        self.resources.insert(uri.into(), text.into());
52        self
53    }
54
55    /// Adds a scripted prompt definition and returns the updated fake server.
56    /// This mutates only the builder's in-memory prompt map; request frames are appended later
57    /// when a client asks the endpoint for prompts.
58    pub fn prompt(mut self, name: impl Into<String>, spec: Value) -> Self {
59        self.prompts.insert(name.into(), spec);
60        self
61    }
62
63    /// Handle next.
64    /// This consumes one queued JSON-RPC frame from the in-memory endpoint and mutates only
65    /// scripted mock state.
66    pub fn handle_next(&mut self, endpoint: &JsonRpcLineEndpoint) -> Result<bool, AgentError> {
67        let frame = match receive_frame_or_parse_error(endpoint)? {
68            ReceiveNext::Empty => return Ok(false),
69            ReceiveNext::Handled => return Ok(true),
70            ReceiveNext::Frame(frame) => frame,
71        };
72        let request = match frame {
73            JsonRpcFrame::Request(request) => request,
74            JsonRpcFrame::Notification(notification) => {
75                if notification.method == "notifications/initialized" {
76                    self.initialized = true;
77                }
78                return Ok(true);
79            }
80            JsonRpcFrame::Response(_) => {
81                endpoint.send_error(None, -32600, "MCP server expected a request frame")?;
82                return Ok(true);
83            }
84        };
85        match request.method.as_str() {
86            "initialize" => {
87                endpoint.send_result(
88                    request.id,
89                    json!({
90                        "protocolVersion": request.params
91                            .get("protocolVersion")
92                            .cloned()
93                            .unwrap_or_else(|| json!("2025-11-25")),
94                        "capabilities": {
95                            "tools": {"listChanged": false},
96                            "resources": {"subscribe": false, "listChanged": false},
97                            "prompts": {"listChanged": false}
98                        },
99                        "serverInfo": {
100                            "name": "agent-sdk-toolkit-mcp-fake",
101                            "version": "0.0.0"
102                        }
103                    }),
104                )?;
105            }
106            "tools/list" => {
107                if !self.ensure_initialized(endpoint, &request)? {
108                    return Ok(true);
109                }
110                endpoint.send_result(request.id, json!({"tools": self.tool_list()}))?;
111            }
112            "tools/call" => {
113                if !self.ensure_initialized(endpoint, &request)? {
114                    return Ok(true);
115                }
116                self.handle_tool_call(endpoint, request)?;
117            }
118            "resources/list" => {
119                if !self.ensure_initialized(endpoint, &request)? {
120                    return Ok(true);
121                }
122                endpoint.send_result(request.id, json!({"resources": self.resource_list()}))?;
123            }
124            "resources/read" => {
125                if !self.ensure_initialized(endpoint, &request)? {
126                    return Ok(true);
127                }
128                self.handle_resource_read(endpoint, request)?;
129            }
130            "prompts/list" => {
131                if !self.ensure_initialized(endpoint, &request)? {
132                    return Ok(true);
133                }
134                endpoint.send_result(request.id, json!({"prompts": self.prompt_list()}))?;
135            }
136            "logging/setLevel" => {
137                if !self.ensure_initialized(endpoint, &request)? {
138                    return Ok(true);
139                }
140                endpoint.send_error(
141                    Some(request.id),
142                    -32010,
143                    "MCP logging control denied by default policy",
144                )?;
145            }
146            _ => {
147                endpoint.send_error(Some(request.id), -32601, "MCP method not found")?;
148            }
149        };
150        Ok(true)
151    }
152
153    /// Request sampling.
154    /// This appends the corresponding JSON-RPC frame to the in-memory test endpoint transcript.
155    pub fn request_sampling(
156        &mut self,
157        endpoint: &JsonRpcLineEndpoint,
158    ) -> Result<JsonRpcId, AgentError> {
159        let id = self.next_client_request_id();
160        endpoint
161            .send_request(
162                id.clone(),
163                "sampling/createMessage",
164                json!({"messages": [{"role": "user", "content": {"type": "text", "text": "sample"}}]}),
165            )
166            .map(|_| id)
167    }
168
169    /// Request elicitation.
170    /// This appends the corresponding JSON-RPC frame to the in-memory test endpoint transcript.
171    pub fn request_elicitation(
172        &mut self,
173        endpoint: &JsonRpcLineEndpoint,
174    ) -> Result<JsonRpcId, AgentError> {
175        let id = self.next_client_request_id();
176        endpoint
177            .send_request(
178                id.clone(),
179                "elicitation/create",
180                json!({"message": "need user input"}),
181            )
182            .map(|_| id)
183    }
184
185    /// Returns the response currently held by this value.
186    /// This reads scripted protocol state or a queued response without contacting an external
187    /// process.
188    pub fn response(&self, endpoint: &JsonRpcLineEndpoint) -> Result<JsonRpcResponse, AgentError> {
189        expect_response(endpoint.receive_frame()?)
190    }
191
192    /// Returns the initialized currently held by this value.
193    /// This reads scripted protocol state or a queued response without contacting an external
194    /// process.
195    pub fn initialized(&self) -> bool {
196        self.initialized
197    }
198
199    fn ensure_initialized(
200        &self,
201        endpoint: &JsonRpcLineEndpoint,
202        request: &JsonRpcRequest,
203    ) -> Result<bool, AgentError> {
204        if self.initialized {
205            return Ok(true);
206        }
207        endpoint.send_error(
208            Some(request.id.clone()),
209            -32002,
210            "MCP client must send notifications/initialized before normal operation",
211        )?;
212        Ok(false)
213    }
214
215    fn next_client_request_id(&mut self) -> JsonRpcId {
216        let id = JsonRpcId::Number(self.next_client_request);
217        self.next_client_request += 1;
218        id
219    }
220
221    fn tool_list(&self) -> Vec<Value> {
222        self.tools
223            .keys()
224            .map(|name| json!({"name": name, "inputSchema": {"type": "object"}}))
225            .collect()
226    }
227
228    fn resource_list(&self) -> Vec<Value> {
229        self.resources
230            .keys()
231            .map(|uri| json!({"uri": uri, "mimeType": "text/plain"}))
232            .collect()
233    }
234
235    fn prompt_list(&self) -> Vec<Value> {
236        self.prompts
237            .iter()
238            .map(|(name, spec)| {
239                let mut spec = spec.clone();
240                if let Some(object) = spec.as_object_mut() {
241                    object.insert("name".to_string(), json!(name));
242                    return spec;
243                }
244                json!({"name": name})
245            })
246            .collect()
247    }
248
249    fn handle_tool_call(
250        &self,
251        endpoint: &JsonRpcLineEndpoint,
252        request: JsonRpcRequest,
253    ) -> Result<(), AgentError> {
254        let name = request
255            .params
256            .get("name")
257            .and_then(Value::as_str)
258            .ok_or_else(|| protocol_violation("MCP tools/call requires name"))?;
259        let Some(result) = self.tools.get(name) else {
260            endpoint.send_error(Some(request.id), -32602, "MCP tool is not available")?;
261            return Ok(());
262        };
263        endpoint.send_result(
264            request.id,
265            json!({
266                "content": [{"type": "text", "text": result.to_string()}],
267                "isError": false
268            }),
269        )?;
270        Ok(())
271    }
272
273    fn handle_resource_read(
274        &self,
275        endpoint: &JsonRpcLineEndpoint,
276        request: JsonRpcRequest,
277    ) -> Result<(), AgentError> {
278        let uri = request
279            .params
280            .get("uri")
281            .and_then(Value::as_str)
282            .ok_or_else(|| protocol_violation("MCP resources/read requires uri"))?;
283        let Some(text) = self.resources.get(uri) else {
284            endpoint.send_error(Some(request.id), -32602, "MCP resource is not available")?;
285            return Ok(());
286        };
287        endpoint.send_result(
288            request.id,
289            json!({
290                "contents": [{
291                    "uri": uri,
292                    "mimeType": "text/plain",
293                    "text": text
294                }]
295            }),
296        )?;
297        Ok(())
298    }
299}
300
301#[derive(Clone, Debug)]
302/// In-memory mcp host proxy fixture for SDK conformance tests.
303/// Use it to script deterministic behavior in memory; any transcript or endpoint mutation is documented on the method that performs it.
304pub struct McpHostProxy {
305    endpoint: JsonRpcLineEndpoint,
306    allowed_tools: BTreeSet<String>,
307    allowed_resources: BTreeSet<String>,
308    allowed_prompts: BTreeSet<String>,
309    pending_methods: BTreeMap<String, String>,
310    next_id: i64,
311}
312
313impl McpHostProxy {
314    /// Creates a new testing::protocol::mcp value with explicit
315    /// caller-provided inputs. This constructor is data-only and
316    /// performs no I/O or external side effects.
317    pub fn new(endpoint: JsonRpcLineEndpoint) -> Self {
318        Self {
319            endpoint,
320            allowed_tools: BTreeSet::new(),
321            allowed_resources: BTreeSet::new(),
322            allowed_prompts: BTreeSet::new(),
323            pending_methods: BTreeMap::new(),
324            next_id: 1,
325        }
326    }
327
328    /// Returns an updated value with allow tool configured.
329    /// This configures or reads the scripted protocol mock in memory.
330    pub fn allow_tool(mut self, name: impl Into<String>) -> Self {
331        self.allowed_tools.insert(name.into());
332        self
333    }
334
335    /// Returns an updated value with allow resource configured.
336    /// This updates the scripted MCP mock allowlist in memory for later protocol calls.
337    pub fn allow_resource(mut self, uri: impl Into<String>) -> Self {
338        self.allowed_resources.insert(uri.into());
339        self
340    }
341
342    /// Returns an updated value with allow prompt configured.
343    /// This updates the scripted MCP mock allowlist in memory for later protocol calls.
344    pub fn allow_prompt(mut self, name: impl Into<String>) -> Self {
345        self.allowed_prompts.insert(name.into());
346        self
347    }
348
349    /// Returns the endpoint currently held by this value.
350    /// This reads scripted protocol state or a queued response without contacting an external
351    /// process.
352    pub fn endpoint(&self) -> &JsonRpcLineEndpoint {
353        &self.endpoint
354    }
355
356    /// Initialize.
357    /// This appends the corresponding JSON-RPC frame to the scripted MCP mock transcript and
358    /// returns the request id when applicable.
359    pub fn initialize(&mut self) -> Result<JsonRpcId, AgentError> {
360        self.request(
361            "initialize",
362            json!({
363                "protocolVersion": "2025-11-25",
364                "capabilities": {},
365                "clientInfo": {
366                    "name": "agent-sdk-toolkit",
367                    "version": "0.0.0"
368                }
369            }),
370        )
371    }
372
373    /// Sends the MCP `notifications/initialized` frame to the scripted server endpoint.
374    /// This appends one in-memory JSON-RPC notification; it does not contact a live MCP process or
375    /// mutate SDK runtime state.
376    pub fn initialized(&self) -> Result<(), AgentError> {
377        self.endpoint
378            .send_notification("notifications/initialized", json!({}))
379            .map(|_| ())
380    }
381
382    /// List tools.
383    /// This appends the corresponding JSON-RPC frame to the in-memory test endpoint transcript.
384    pub fn list_tools(&mut self) -> Result<JsonRpcId, AgentError> {
385        self.request("tools/list", json!({}))
386    }
387
388    /// List resources.
389    /// This appends the corresponding JSON-RPC frame to the scripted MCP mock transcript and
390    /// returns the request id when applicable.
391    pub fn list_resources(&mut self) -> Result<JsonRpcId, AgentError> {
392        self.request("resources/list", json!({}))
393    }
394
395    /// List prompts.
396    /// This appends the corresponding JSON-RPC frame to the scripted MCP mock transcript and
397    /// returns the request id when applicable.
398    pub fn list_prompts(&mut self) -> Result<JsonRpcId, AgentError> {
399        self.request("prompts/list", json!({}))
400    }
401
402    /// Call tool.
403    /// This appends the corresponding JSON-RPC frame to the in-memory test endpoint transcript.
404    pub fn call_tool(&mut self, name: &str, arguments: Value) -> Result<JsonRpcId, AgentError> {
405        if !self.allowed_tools.contains(name) {
406            return Err(policy_denial(format!(
407                "MCP tool {name} is not selected in host proxy policy"
408            )));
409        }
410        self.request("tools/call", json!({"name": name, "arguments": arguments}))
411    }
412
413    /// Read resource.
414    /// This appends the corresponding JSON-RPC frame to the in-memory test endpoint transcript.
415    pub fn read_resource(&mut self, uri: &str) -> Result<JsonRpcId, AgentError> {
416        if !self.allowed_resources.contains(uri) {
417            return Err(policy_denial(format!(
418                "MCP resource {uri} is not selected in host proxy policy"
419            )));
420        }
421        self.request("resources/read", json!({"uri": uri}))
422    }
423
424    /// Handle next.
425    /// This consumes one queued JSON-RPC frame from the endpoint and mutates only the scripted
426    /// mock state.
427    pub fn handle_next(&self, endpoint: &JsonRpcLineEndpoint) -> Result<bool, AgentError> {
428        let frame = match receive_frame_or_parse_error(endpoint)? {
429            ReceiveNext::Empty => return Ok(false),
430            ReceiveNext::Handled => return Ok(true),
431            ReceiveNext::Frame(frame) => frame,
432        };
433        let JsonRpcFrame::Request(request) = frame else {
434            return Ok(true);
435        };
436        match request.method.as_str() {
437            "sampling/createMessage" | "elicitation/create" => endpoint.send_error(
438                Some(request.id),
439                -32010,
440                "MCP server-to-client request denied by host proxy policy",
441            )?,
442            _ => endpoint.send_error(Some(request.id), -32601, "MCP client method not found")?,
443        };
444        Ok(true)
445    }
446
447    /// Returns the response currently held by this value.
448    /// This reads scripted protocol state or a queued response without contacting an external
449    /// process.
450    pub fn response(&mut self) -> Result<JsonRpcResponse, AgentError> {
451        let mut response = expect_response(self.endpoint.receive_frame()?)?;
452        if response.id == JsonRpcId::Null {
453            return Ok(response);
454        }
455        let Some(method) = self.pending_methods.remove(&response.id.as_key()) else {
456            return Err(protocol_violation("unexpected MCP response id"));
457        };
458        self.apply_response_policy(&method, &mut response);
459        Ok(response)
460    }
461
462    /// Returns allowed tool names from response for the current value.
463    /// This is a read-only or data-construction helper unless the method body explicitly calls
464    /// a port or store.
465    pub fn allowed_tool_names_from_response(
466        &self,
467        response: &JsonRpcResponse,
468    ) -> Result<Vec<String>, AgentError> {
469        let tools = response
470            .result
471            .as_ref()
472            .and_then(|value| value.get("tools"))
473            .and_then(Value::as_array)
474            .ok_or_else(|| protocol_violation("MCP tools/list response missing tools array"))?;
475        let mut names = tools
476            .iter()
477            .filter_map(|tool| tool.get("name").and_then(Value::as_str))
478            .filter(|name| self.allowed_tools.contains(*name))
479            .map(str::to_string)
480            .collect::<Vec<_>>();
481        names.sort();
482        Ok(names)
483    }
484
485    fn request(&mut self, method: &str, params: Value) -> Result<JsonRpcId, AgentError> {
486        let id = JsonRpcId::Number(self.next_id);
487        self.next_id += 1;
488        self.endpoint.send_request(id.clone(), method, params)?;
489        self.pending_methods.insert(id.as_key(), method.to_string());
490        Ok(id)
491    }
492
493    fn apply_response_policy(&self, method: &str, response: &mut JsonRpcResponse) {
494        if response.error.is_some() {
495            return;
496        }
497        let Some(result) = response.result.as_mut() else {
498            return;
499        };
500        match method {
501            "tools/list" => filter_named_array(result, "tools", "name", &self.allowed_tools),
502            "resources/list" => {
503                filter_named_array(result, "resources", "uri", &self.allowed_resources)
504            }
505            "prompts/list" => filter_named_array(result, "prompts", "name", &self.allowed_prompts),
506            _ => {}
507        }
508    }
509}
510
511fn receive_frame_or_parse_error(endpoint: &JsonRpcLineEndpoint) -> Result<ReceiveNext, AgentError> {
512    let Some(line) = endpoint.try_receive_raw_line()? else {
513        return Ok(ReceiveNext::Empty);
514    };
515    match JsonRpcFrame::from_line(&line) {
516        Ok(frame) => Ok(ReceiveNext::Frame(frame)),
517        Err(error) => {
518            endpoint.send_error(None, -32700, error.context().message)?;
519            Ok(ReceiveNext::Handled)
520        }
521    }
522}
523
524fn filter_named_array(result: &mut Value, field: &str, key: &str, allowed: &BTreeSet<String>) {
525    let Some(items) = result.get_mut(field).and_then(Value::as_array_mut) else {
526        return;
527    };
528    items.retain(|item| {
529        item.get(key)
530            .and_then(Value::as_str)
531            .is_some_and(|name| allowed.contains(name))
532    });
533}
534
535fn policy_denial(message: impl Into<String>) -> AgentError {
536    AgentError::new(
537        AgentErrorKind::PolicyDenial,
538        RetryClassification::UserActionNeeded,
539        message,
540    )
541}
agent_sdk_toolkit/testing/protocol/mcp.rs

agent_sdk_toolkit/testing/protocol/
mcp.rs
