Skip to main content

max / makenotwork

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