Skip to main content

max / makenotwork

13.9 KB · 482 lines History Blame Raw
1 //! End-to-end tests for the Streamable HTTP MCP server: drive a real bound
2 //! server over HTTP with a plain `reqwest` client (no `Agent`), exercising
3 //! `initialize`, `tools/list` (both projections), and `tools/call` across the
4 //! read / write-granted / write-denied / refusal / not-found paths.
5
6 #![cfg(feature = "server")]
7
8 use std::collections::HashSet;
9 use std::time::Duration;
10
11 use async_trait::async_trait;
12 use kberg::server::{self, ServeConfig};
13 use kberg::{
14 Refusal, ResourceContents, ResourceDescriptor, ResourceRegistry, Result, SurfaceProjection,
15 Tool, ToolCallResult, ToolKind, ToolRegistry, WriteCapability,
16 };
17 use serde_json::{Value, json};
18
19 struct Ping;
20 #[async_trait]
21 impl Tool for Ping {
22 fn name(&self) -> &'static str {
23 "ping"
24 }
25 fn description(&self) -> &'static str {
26 "read-only ping"
27 }
28 fn kind(&self) -> ToolKind {
29 ToolKind::Read
30 }
31 fn small_model_safe(&self) -> bool {
32 true
33 }
34 fn input_schema(&self) -> Value {
35 json!({ "type": "object", "properties": {} })
36 }
37 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
38 Ok(ToolCallResult::text("pong"))
39 }
40 }
41
42 struct Save;
43 #[async_trait]
44 impl Tool for Save {
45 fn name(&self) -> &'static str {
46 "save"
47 }
48 fn description(&self) -> &'static str {
49 "write, gated on `demo.save`"
50 }
51 fn kind(&self) -> ToolKind {
52 ToolKind::Write(WriteCapability::new("demo.save", "persist a value"))
53 }
54 fn input_schema(&self) -> Value {
55 json!({ "type": "object", "properties": { "v": { "type": "string" } } })
56 }
57 async fn call(&self, args: Value) -> Result<ToolCallResult> {
58 let v = args.get("v").and_then(Value::as_str).unwrap_or("");
59 Ok(ToolCallResult::text(format!("saved {v}")))
60 }
61 }
62
63 /// reqwest is declared `rustls-no-provider` (see the manifest: the default
64 /// feature set welds aws-lc-rs in), so a client cannot be built until a
65 /// provider is installed. Idempotent, so every test calls it.
66 fn client() -> reqwest::Client {
67 let _ = rustls::crypto::ring::default_provider().install_default();
68 reqwest::Client::new()
69 }
70
71 fn registry() -> ToolRegistry {
72 let mut r = ToolRegistry::new();
73 r.register(Ping);
74 r.register(Save);
75 r.register(Refusal::new("nuke", "not offered", "never exposed"));
76 r
77 }
78
79 /// A read tool that never returns on its own — only cancellation ends it.
80 struct Hang;
81 #[async_trait]
82 impl Tool for Hang {
83 fn name(&self) -> &'static str {
84 "hang"
85 }
86 fn description(&self) -> &'static str {
87 "read-only; blocks until cancelled"
88 }
89 fn kind(&self) -> ToolKind {
90 ToolKind::Read
91 }
92 fn input_schema(&self) -> Value {
93 json!({ "type": "object", "properties": {} })
94 }
95 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
96 std::future::pending::<()>().await;
97 unreachable!()
98 }
99 }
100
101 async fn rpc(client: &reqwest::Client, url: &str, method: &str, params: Value) -> Value {
102 client
103 .post(url)
104 .json(&json!({ "jsonrpc": "2.0", "id": 1, "method": method, "params": params }))
105 .send()
106 .await
107 .unwrap()
108 .json()
109 .await
110 .unwrap()
111 }
112
113 /// POST a request under a session id and return the parsed JSON-RPC response.
114 async fn rpc_in_session(
115 client: &reqwest::Client,
116 url: &str,
117 session: &str,
118 id: i64,
119 method: &str,
120 params: Value,
121 ) -> Value {
122 client
123 .post(url)
124 .header("Mcp-Session-Id", session)
125 .json(&json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }))
126 .send()
127 .await
128 .unwrap()
129 .json()
130 .await
131 .unwrap()
132 }
133
134 /// Initialize and return the server-assigned `Mcp-Session-Id`.
135 async fn open_session(client: &reqwest::Client, url: &str) -> String {
136 let resp = client
137 .post(url)
138 .json(&json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize" }))
139 .send()
140 .await
141 .unwrap();
142 resp.headers()
143 .get("Mcp-Session-Id")
144 .expect("initialize returns a session id")
145 .to_str()
146 .unwrap()
147 .to_string()
148 }
149
150 /// Read the next `data:` JSON payload off an SSE response, buffering across
151 /// chunks. Returns `None` at end of stream.
152 async fn next_sse_json(resp: &mut reqwest::Response, buf: &mut String) -> Option<Value> {
153 loop {
154 if let Some(pos) = buf.find("data:")
155 && let Some(nl) = buf[pos..].find('\n')
156 {
157 let line = buf[pos + 5..pos + nl].trim().to_string();
158 *buf = buf[pos + nl + 1..].to_string();
159 if !line.is_empty()
160 && let Ok(v) = serde_json::from_str::<Value>(&line)
161 {
162 return Some(v);
163 }
164 continue;
165 }
166 match resp.chunk().await.ok()? {
167 Some(bytes) => buf.push_str(&String::from_utf8_lossy(&bytes)),
168 None => return None,
169 }
170 }
171 }
172
173 #[tokio::test]
174 async fn full_surface_round_trip() {
175 let bound = server::serve(
176 registry(),
177 "127.0.0.1:0",
178 ServeConfig::granting(["demo.save"]),
179 )
180 .await
181 .unwrap();
182 let url = bound.url();
183 let client = client();
184
185 // initialize
186 let init = rpc(&client, &url, "initialize", Value::Null).await;
187 assert_eq!(init["result"]["serverInfo"]["name"], "kberg");
188 assert!(init["result"]["protocolVersion"].is_string());
189
190 // tools/list — full surface has all three, with kind/refusal extensions
191 let list = rpc(&client, &url, "tools/list", Value::Null).await;
192 let tools = list["result"]["tools"].as_array().unwrap();
193 assert_eq!(tools.len(), 3);
194 let save = tools.iter().find(|t| t["name"] == "save").unwrap();
195 assert_eq!(save["kbergKind"], "write");
196 assert_eq!(save["kbergCapability"]["id"], "demo.save");
197 let nuke = tools.iter().find(|t| t["name"] == "nuke").unwrap();
198 assert!(nuke["kbergRefusalReason"].is_string());
199
200 // tools/call — read works
201 let ping = rpc(
202 &client,
203 &url,
204 "tools/call",
205 json!({ "name": "ping", "arguments": {} }),
206 )
207 .await;
208 assert_eq!(ping["result"]["content"][0]["text"], "pong");
209
210 // tools/call — granted write works
211 let save = rpc(
212 &client,
213 &url,
214 "tools/call",
215 json!({ "name": "save", "arguments": { "v": "x" } }),
216 )
217 .await;
218 assert_eq!(save["result"]["content"][0]["text"], "saved x");
219
220 // tools/call — refusal returns a JSON-RPC error, not a tool result
221 let nuke = rpc(
222 &client,
223 &url,
224 "tools/call",
225 json!({ "name": "nuke", "arguments": {} }),
226 )
227 .await;
228 assert!(nuke["error"].is_object());
229
230 // tools/call — unknown tool errors
231 let miss = rpc(
232 &client,
233 &url,
234 "tools/call",
235 json!({ "name": "ghost", "arguments": {} }),
236 )
237 .await;
238 assert!(miss["error"].is_object());
239
240 bound.shutdown();
241 }
242
243 #[tokio::test]
244 async fn ungranted_write_is_denied_over_the_wire() {
245 // Server started with an empty grant set: `save` must be refused.
246 let bound = server::serve(
247 registry(),
248 "127.0.0.1:0",
249 ServeConfig::granting(Vec::<String>::new()),
250 )
251 .await
252 .unwrap();
253 let url = bound.url();
254 let client = client();
255
256 let save = rpc(
257 &client,
258 &url,
259 "tools/call",
260 json!({ "name": "save", "arguments": { "v": "x" } }),
261 )
262 .await;
263 let msg = save["error"]["message"].as_str().unwrap();
264 assert!(
265 msg.contains("demo.save"),
266 "denial names the capability: {msg}"
267 );
268 assert!(
269 msg.contains("do not retry"),
270 "denial tells small models not to retry: {msg}"
271 );
272
273 bound.shutdown();
274 }
275
276 #[tokio::test]
277 async fn resources_listed_read_and_removed_over_the_wire() {
278 // A resource registry the app mutates at runtime, mirroring deox opening
279 // and closing buffers while the server stays up.
280 let resources = ResourceRegistry::new();
281 resources.register(
282 ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
283 |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
284 );
285
286 let bound = server::serve(
287 registry(),
288 "127.0.0.1:0",
289 ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
290 )
291 .await
292 .unwrap();
293 let url = bound.url();
294 let client = client();
295
296 // initialize advertises the resources capability with subscribe support.
297 let init = rpc(&client, &url, "initialize", Value::Null).await;
298 assert_eq!(
299 init["result"]["capabilities"]["resources"]["subscribe"],
300 true
301 );
302
303 // resources/list
304 let list = rpc(&client, &url, "resources/list", Value::Null).await;
305 let items = list["result"]["resources"].as_array().unwrap();
306 assert_eq!(items.len(), 1);
307 assert_eq!(items[0]["uri"], "buffer://main.rs");
308 assert_eq!(items[0]["mimeType"], "text/x-rust");
309
310 // resources/read
311 let read = rpc(
312 &client,
313 &url,
314 "resources/read",
315 json!({ "uri": "buffer://main.rs" }),
316 )
317 .await;
318 assert_eq!(read["result"]["contents"][0]["text"], "fn main() {}");
319
320 // Close the buffer at runtime; the running server observes it (shared Arc).
321 assert!(resources.remove("buffer://main.rs"));
322 let gone = rpc(
323 &client,
324 &url,
325 "resources/read",
326 json!({ "uri": "buffer://main.rs" }),
327 )
328 .await;
329 assert_eq!(gone["error"]["code"], -32002); // RESOURCE_NOT_FOUND
330
331 bound.shutdown();
332 }
333
334 #[tokio::test]
335 async fn compact_projection_hides_write_tools() {
336 let bound = server::serve(
337 registry(),
338 "127.0.0.1:0",
339 ServeConfig::granting(["demo.save"]).with_projection(SurfaceProjection::Compact),
340 )
341 .await
342 .unwrap();
343 let url = bound.url();
344 let client = client();
345
346 let list = rpc(&client, &url, "tools/list", Value::Null).await;
347 let names: HashSet<String> = list["result"]["tools"]
348 .as_array()
349 .unwrap()
350 .iter()
351 .map(|t| t["name"].as_str().unwrap().to_string())
352 .collect();
353 assert!(names.contains("ping"));
354 assert!(!names.contains("save"));
355 assert!(!names.contains("nuke"));
356
357 bound.shutdown();
358 }
359
360 #[tokio::test]
361 async fn resource_update_is_pushed_over_sse() {
362 // The deox flow: agent subscribes to a buffer, the human edits it, the
363 // agent's SSE stream receives a `resources/updated` push.
364 let resources = ResourceRegistry::new();
365 resources.register(
366 ResourceDescriptor::new("buffer://main.rs", "main.rs"),
367 |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
368 );
369 let bound = server::serve(
370 registry(),
371 "127.0.0.1:0",
372 ServeConfig::granting(Vec::<String>::new()).with_resources(resources.clone()),
373 )
374 .await
375 .unwrap();
376 let url = bound.url();
377 let client = client();
378
379 let session = open_session(&client, &url).await;
380
381 // Open the server->client SSE stream (registers the update forwarder).
382 let mut sse = client
383 .get(&url)
384 .header("Mcp-Session-Id", &session)
385 .send()
386 .await
387 .unwrap();
388 let mut buf = String::new();
389
390 // Subscribe to the buffer.
391 let sub = rpc_in_session(
392 &client,
393 &url,
394 &session,
395 2,
396 "resources/subscribe",
397 json!({ "uri": "buffer://main.rs" }),
398 )
399 .await;
400 assert!(
401 sub["result"].is_object(),
402 "subscribe returns an empty result: {sub}"
403 );
404
405 // The app mutates the buffer and announces it.
406 resources.notify_updated("buffer://main.rs");
407
408 let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
409 .await
410 .expect("an update arrives before the timeout")
411 .expect("the stream yields a notification");
412 assert_eq!(note["method"], "notifications/resources/updated");
413 assert_eq!(note["params"]["uri"], "buffer://main.rs");
414
415 bound.shutdown();
416 }
417
418 #[tokio::test]
419 async fn app_cancel_aborts_call_and_notifies_over_sse() {
420 // The deox cancel: a call is in flight, the human cancels, the server
421 // aborts the call and tells the agent over SSE.
422 let mut reg = ToolRegistry::new();
423 reg.register(Hang);
424 let bound = server::serve(
425 reg,
426 "127.0.0.1:0",
427 ServeConfig::granting(Vec::<String>::new()),
428 )
429 .await
430 .unwrap();
431 let url = bound.url();
432 let client = client();
433
434 let session = open_session(&client, &url).await;
435 let mut sse = client
436 .get(&url)
437 .header("Mcp-Session-Id", &session)
438 .send()
439 .await
440 .unwrap();
441 let mut buf = String::new();
442
443 // Fire a call that hangs until cancelled.
444 let call = {
445 let client = client.clone();
446 let url = url.clone();
447 let session = session.clone();
448 tokio::spawn(async move {
449 rpc_in_session(
450 &client,
451 &url,
452 &session,
453 7,
454 "tools/call",
455 json!({ "name": "hang", "arguments": {} }),
456 )
457 .await
458 })
459 };
460
461 // Let it register as in-flight, then cancel it app-side.
462 tokio::time::sleep(Duration::from_millis(200)).await;
463 assert_eq!(bound.cancel_all("human took the buffer back"), 1);
464
465 // The agent's SSE stream receives a cancelled notification echoing id 7.
466 let note = tokio::time::timeout(Duration::from_secs(2), next_sse_json(&mut sse, &mut buf))
467 .await
468 .expect("a cancelled notification arrives")
469 .expect("the stream yields a notification");
470 assert_eq!(note["method"], "notifications/cancelled");
471 assert_eq!(note["params"]["requestId"], 7);
472
473 // The blocked POST completes with a request-cancelled JSON-RPC error.
474 let resp = tokio::time::timeout(Duration::from_secs(2), call)
475 .await
476 .expect("the call returns after cancellation")
477 .unwrap();
478 assert_eq!(resp["error"]["code"], -32800);
479
480 bound.shutdown();
481 }
482