Skip to main content

max / makenotwork

23.3 KB · 701 lines History Blame Raw
1 //! Streamable HTTP MCP server.
2 //!
3 //! One `/mcp` route, three methods:
4 //! - **POST** — JSON-RPC requests: `initialize` (opens a session, returns its
5 //! id in `Mcp-Session-Id`), `tools/list`, `tools/call`, `resources/list`,
6 //! `resources/read`, `resources/subscribe`/`unsubscribe`, and the
7 //! `notifications/cancelled` notification.
8 //! - **GET** — opens the server→client SSE stream for a session, down which
9 //! `notifications/resources/updated` and `notifications/cancelled` are
10 //! pushed.
11 //! - **DELETE** — ends a session.
12 //!
13 //! Configuration is passed at startup via [`ServeConfig`]:
14 //! - `projection` — Full or Compact surface exposed to clients on this port.
15 //! - `grants` — which write capabilities are honored; unmatched writes return
16 //! `CapabilityDenied`. `None` bypasses the check (fully-trusted context).
17 //! - `resources` — an optional [`ResourceRegistry`]; when set, the resource
18 //! methods above are served and the `resources` capability is advertised.
19
20 use std::collections::HashSet;
21 use std::convert::Infallible;
22 use std::net::SocketAddr;
23 use std::sync::Arc;
24
25 use axum::{
26 Json, Router,
27 extract::State,
28 http::{HeaderMap, StatusCode},
29 response::{
30 IntoResponse,
31 sse::{Event, KeepAlive, Sse},
32 },
33 routing::post,
34 };
35 use serde_json::{Value, json};
36 use tokio::net::TcpListener;
37 use tokio::task::JoinHandle;
38 use tokio_stream::wrappers::UnboundedReceiverStream;
39 use tokio_stream::{Stream, StreamExt};
40
41 use crate::error::{Error, Result};
42 use crate::resource::{ReadResourceParams, ResourceRegistry, read_result};
43 use crate::tool::{SurfaceProjection, ToolRegistry};
44
45 pub mod protocol;
46 mod session;
47 pub mod stdio;
48
49 use protocol::{
50 CallToolParams, CancelledParams, InitializeResult, JsonRpcRequest, JsonRpcResponse,
51 MCP_PROTOCOL_VERSION, ResourcesCapability, ServerCapabilities, ServerInfo, SubscribeParams,
52 ToolsCapability, codes,
53 };
54 use session::{Session, Sessions, resource_updated_notification};
55
56 /// Header carrying the MCP session id, per the Streamable HTTP transport.
57 const SESSION_HEADER: &str = "mcp-session-id";
58
59 pub struct ServeConfig {
60 pub projection: SurfaceProjection,
61 pub grants: Option<HashSet<String>>,
62 /// Resource surface exposed alongside the tools. `None` means the server
63 /// advertises no `resources` capability and rejects `resources/*` methods.
64 pub resources: Option<ResourceRegistry>,
65 }
66
67 impl Default for ServeConfig {
68 fn default() -> Self {
69 Self {
70 projection: SurfaceProjection::Full,
71 grants: Some(HashSet::new()),
72 resources: None,
73 }
74 }
75 }
76
77 impl ServeConfig {
78 /// Full surface, honoring exactly the given write-capability ids. The
79 /// common shape for a standalone binary that knows which capabilities its
80 /// driving session should hold.
81 pub fn granting(ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
82 Self {
83 projection: SurfaceProjection::Full,
84 grants: Some(ids.into_iter().map(Into::into).collect()),
85 resources: None,
86 }
87 }
88
89 #[must_use]
90 pub fn with_projection(mut self, projection: SurfaceProjection) -> Self {
91 self.projection = projection;
92 self
93 }
94
95 /// Attach a resource registry, enabling `resources/list` and
96 /// `resources/read` and advertising the `resources` capability.
97 #[must_use]
98 pub fn with_resources(mut self, resources: ResourceRegistry) -> Self {
99 self.resources = Some(resources);
100 self
101 }
102 }
103
104 /// Handle to a running MCP server.
105 ///
106 /// Dropping this does *not* stop the server — call [`BoundServer::shutdown`]
107 /// or abort the task explicitly.
108 pub struct BoundServer {
109 pub addr: SocketAddr,
110 handle: JoinHandle<std::io::Result<()>>,
111 sessions: Sessions,
112 }
113
114 impl BoundServer {
115 /// Full URL clients POST to, e.g. `http://127.0.0.1:54321/mcp`.
116 pub fn url(&self) -> String {
117 format!("http://{}/mcp", self.addr)
118 }
119
120 /// Cancel every in-flight `tools/call` across all sessions, pushing a
121 /// `notifications/cancelled` to each affected client. The app-initiated
122 /// path — a human taking a buffer back in deox. Returns the count
123 /// cancelled. For v0 (one agent = one session) this is the whole cancel
124 /// surface; [`cancel_session`](Self::cancel_session) targets one client
125 /// once an app tracks session ids.
126 pub fn cancel_all(&self, reason: &str) -> usize {
127 self.sessions.cancel_all(reason)
128 }
129
130 /// Cancel in-flight `tools/call`s for one session by its `Mcp-Session-Id`.
131 /// Returns the count cancelled (0 if the session is unknown or idle).
132 pub fn cancel_session(&self, session_id: &str, reason: &str) -> usize {
133 self.sessions
134 .get(session_id)
135 .map_or(0, |s| s.cancel_all(reason))
136 }
137
138 /// Ids of the currently-connected sessions.
139 pub fn sessions(&self) -> Vec<String> {
140 self.sessions.ids()
141 }
142
143 pub fn shutdown(self) {
144 self.handle.abort();
145 }
146
147 /// Block until the server task finishes.
148 ///
149 /// A standalone MCP binary calls this after [`serve`] to keep the process
150 /// alive for as long as the server is listening. The task only ends on an
151 /// unrecoverable listener error (or an external abort), so under normal
152 /// operation this future never resolves — it is the binary's park point.
153 pub async fn wait(self) -> Result<()> {
154 match self.handle.await {
155 Ok(io_result) => io_result.map_err(Error::from),
156 Err(join_err) if join_err.is_cancelled() => Ok(()),
157 Err(join_err) => Err(Error::Transport(join_err.to_string())),
158 }
159 }
160 }
161
162 #[derive(Clone)]
163 struct AppState {
164 registry: ToolRegistry,
165 config: Arc<ServeConfig>,
166 sessions: Sessions,
167 }
168
169 /// Start an MCP server. Returns once the listener is bound; the server runs
170 /// in a background tokio task.
171 pub async fn serve(registry: ToolRegistry, bind: &str, config: ServeConfig) -> Result<BoundServer> {
172 let listener = TcpListener::bind(bind).await?;
173 let addr = listener.local_addr()?;
174
175 let sessions = Sessions::default();
176 let state = AppState {
177 registry,
178 config: Arc::new(config),
179 sessions: sessions.clone(),
180 };
181
182 let app = Router::new()
183 .route(
184 "/mcp",
185 post(handle_rpc).get(handle_sse).delete(handle_delete),
186 )
187 .with_state(state);
188
189 let handle = tokio::spawn(async move { axum::serve(listener, app).await });
190
191 tracing::info!(%addr, "kberg server listening");
192 Ok(BoundServer {
193 addr,
194 handle,
195 sessions,
196 })
197 }
198
199 /// Look up the session named by the `Mcp-Session-Id` header, if any.
200 fn session_from_headers(state: &AppState, headers: &HeaderMap) -> Option<Arc<Session>> {
201 headers
202 .get(SESSION_HEADER)
203 .and_then(|h| h.to_str().ok())
204 .and_then(|id| state.sessions.get(id))
205 }
206
207 async fn handle_rpc(
208 State(state): State<AppState>,
209 headers: HeaderMap,
210 Json(req): Json<Value>,
211 ) -> impl IntoResponse {
212 let req: JsonRpcRequest = match serde_json::from_value(req) {
213 Ok(r) => r,
214 Err(e) => {
215 return (
216 StatusCode::BAD_REQUEST,
217 Json(JsonRpcResponse::err(
218 Value::Null,
219 codes::PARSE_ERROR,
220 e.to_string(),
221 )),
222 )
223 .into_response();
224 }
225 };
226
227 // `initialize` opens a session; every other method rides an existing one
228 // named by the header (absent for a bare client, e.g. over stdio-style use).
229 let is_initialize = req.method == "initialize";
230 let session = session_from_headers(&state, &headers);
231
232 let dispatched = dispatch(
233 &state.registry,
234 state.config.resources.as_ref(),
235 state.config.projection,
236 state.config.grants.as_ref(),
237 session.as_deref(),
238 req,
239 )
240 .await;
241
242 match dispatched {
243 // Notifications (no `id`) get a 202 with no body.
244 None => StatusCode::ACCEPTED.into_response(),
245 Some(response) => {
246 let mut http = (StatusCode::OK, Json(response)).into_response();
247 if is_initialize {
248 let new_session = state.sessions.create();
249 if let Ok(value) = new_session.id.parse() {
250 http.headers_mut().insert(SESSION_HEADER, value);
251 }
252 }
253 http
254 }
255 }
256 }
257
258 /// GET `/mcp` — open the server→client SSE stream for a session. This is the
259 /// channel `notifications/resources/updated` and `notifications/cancelled`
260 /// travel down. Requires a valid `Mcp-Session-Id`.
261 async fn handle_sse(State(state): State<AppState>, headers: HeaderMap) -> axum::response::Response {
262 let Some(session) = session_from_headers(&state, &headers) else {
263 return (StatusCode::BAD_REQUEST, "missing or unknown Mcp-Session-Id").into_response();
264 };
265
266 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
267 session.attach_out(tx.clone());
268
269 // Forward resource updates this session is subscribed to. The task lives as
270 // long as the SSE stream: once the client disconnects, `rx` drops and the
271 // `tx.send` below fails, ending the loop.
272 if let Some(resources) = &state.config.resources {
273 let mut updates = resources.subscribe_updates();
274 let session = session.clone();
275 tokio::spawn(async move {
276 loop {
277 match updates.recv().await {
278 Ok(uri) => {
279 if session.is_subscribed(&uri)
280 && tx.send(resource_updated_notification(&uri)).is_err()
281 {
282 break; // client gone
283 }
284 }
285 // A slow session that lagged past the buffer drops those
286 // updates and keeps streaming; it re-reads on next use.
287 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
288 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
289 }
290 }
291 });
292 }
293
294 let stream = UnboundedReceiverStream::new(rx)
295 .map(|note| Ok::<Event, Infallible>(Event::default().data(note.to_string())));
296 sse_response(stream)
297 }
298
299 fn sse_response<S>(stream: S) -> axum::response::Response
300 where
301 S: Stream<Item = std::result::Result<Event, Infallible>> + Send + 'static,
302 {
303 Sse::new(stream)
304 .keep_alive(KeepAlive::default())
305 .into_response()
306 }
307
308 /// DELETE `/mcp` — end a session (the client is done). Idempotent.
309 async fn handle_delete(State(state): State<AppState>, headers: HeaderMap) -> impl IntoResponse {
310 if let Some(id) = headers.get(SESSION_HEADER).and_then(|h| h.to_str().ok()) {
311 state.sessions.remove(id);
312 }
313 StatusCode::NO_CONTENT
314 }
315
316 /// Transport-agnostic JSON-RPC dispatch. Handles `initialize`, `tools/list`,
317 /// and `tools/call` against a registry. Returns `None` for a notification
318 /// (a request with no `id`), which carries no response. Shared by the HTTP
319 /// handler and the [`stdio`] transport.
320 pub(crate) async fn dispatch(
321 registry: &ToolRegistry,
322 resources: Option<&ResourceRegistry>,
323 projection: SurfaceProjection,
324 grants: Option<&HashSet<String>>,
325 session: Option<&Session>,
326 req: JsonRpcRequest,
327 ) -> Option<JsonRpcResponse> {
328 // Notifications (no `id`) carry no response, but some carry a side effect.
329 let Some(id) = req.id.clone() else {
330 if req.method == "notifications/cancelled"
331 && let Some(session) = session
332 && let Ok(params) =
333 serde_json::from_value::<CancelledParams>(req.params.unwrap_or(Value::Null))
334 {
335 session.cancel_request(&params.request_id.to_string());
336 }
337 return None;
338 };
339
340 let response = match req.method.as_str() {
341 "initialize" => JsonRpcResponse::ok(id, initialize_result(resources.is_some())),
342 "tools/list" => {
343 let specs = registry.specs_projected(projection);
344 JsonRpcResponse::ok(id, json!({ "tools": specs }))
345 }
346 "tools/call" => {
347 match serde_json::from_value::<CallToolParams>(req.params.unwrap_or(Value::Null)) {
348 Ok(params) => {
349 let call = registry.call(&params.name, params.arguments, grants);
350 // With a session, run so a human/client can cancel mid-flight.
351 let result = match session {
352 Some(session) => session.run_cancellable(&id, call).await,
353 None => call.await,
354 };
355 match result {
356 Ok(r) => JsonRpcResponse::ok(id, serde_json::to_value(r).unwrap()),
357 Err(e) => rpc_err_from(id, e),
358 }
359 }
360 Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
361 }
362 }
363 "resources/list" => match resources {
364 Some(res) => JsonRpcResponse::ok(id, json!({ "resources": res.list() })),
365 None => resources_unsupported(id),
366 },
367 "resources/read" => match resources {
368 Some(res) => match serde_json::from_value::<ReadResourceParams>(
369 req.params.unwrap_or(Value::Null),
370 ) {
371 Ok(params) => match res.read(&params.uri).await {
372 Ok(contents) => JsonRpcResponse::ok(id, read_result(contents)),
373 Err(e) => rpc_err_from(id, e),
374 },
375 Err(e) => JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
376 },
377 None => resources_unsupported(id),
378 },
379 method @ ("resources/subscribe" | "resources/unsubscribe") => {
380 dispatch_subscription(method, resources, session, id, req.params)
381 }
382 other => JsonRpcResponse::err(
383 id,
384 codes::METHOD_NOT_FOUND,
385 format!("unknown method: {other}"),
386 ),
387 };
388 Some(response)
389 }
390
391 /// `resources/subscribe` / `resources/unsubscribe`. Both need a resource
392 /// registry (to validate the uri exists) and a session (to hold the
393 /// subscription) — the latter means these are HTTP-only.
394 fn dispatch_subscription(
395 method: &str,
396 resources: Option<&ResourceRegistry>,
397 session: Option<&Session>,
398 id: Value,
399 params: Option<Value>,
400 ) -> JsonRpcResponse {
401 let Some(resources) = resources else {
402 return resources_unsupported(id);
403 };
404 let Some(session) = session else {
405 return JsonRpcResponse::err(
406 id,
407 codes::INVALID_PARAMS,
408 "resource subscriptions require a session (HTTP transport)",
409 );
410 };
411 let params: SubscribeParams = match serde_json::from_value(params.unwrap_or(Value::Null)) {
412 Ok(p) => p,
413 Err(e) => return JsonRpcResponse::err(id, codes::INVALID_PARAMS, e.to_string()),
414 };
415 if method == "resources/subscribe" {
416 if !resources.contains(&params.uri) {
417 return rpc_err_from(id, Error::ResourceNotFound(params.uri));
418 }
419 session.subscribe(params.uri);
420 } else {
421 session.unsubscribe(&params.uri);
422 }
423 // MCP subscribe/unsubscribe return an empty result on success.
424 JsonRpcResponse::ok(id, json!({}))
425 }
426
427 /// A `resources/*` method arrived at a server with no resource registry
428 /// configured. It is a method the server does not implement, not a bad request.
429 fn resources_unsupported(id: Value) -> JsonRpcResponse {
430 JsonRpcResponse::err(id, codes::METHOD_NOT_FOUND, "server exposes no resources")
431 }
432
433 // id/e are consumed into the returned response.
434 #[allow(clippy::needless_pass_by_value)]
435 fn rpc_err_from(id: Value, e: Error) -> JsonRpcResponse {
436 let (code, message) = match &e {
437 Error::ToolNotFound(_) => (codes::METHOD_NOT_FOUND, e.to_string()),
438 Error::ResourceNotFound(_) => (codes::RESOURCE_NOT_FOUND, e.to_string()),
439 Error::Cancelled(_) => (codes::REQUEST_CANCELLED, e.to_string()),
440 Error::InvalidArgs { .. } => (codes::INVALID_PARAMS, e.to_string()),
441 Error::CapabilityDenied { .. } | Error::Refused { .. } => {
442 (codes::INVALID_PARAMS, e.to_string())
443 }
444 _ => (codes::INTERNAL_ERROR, e.to_string()),
445 };
446 JsonRpcResponse::err(id, code, message)
447 }
448
449 fn initialize_result(resources_enabled: bool) -> Value {
450 serde_json::to_value(InitializeResult {
451 protocol_version: MCP_PROTOCOL_VERSION,
452 capabilities: ServerCapabilities {
453 tools: ToolsCapability {
454 list_changed: false,
455 },
456 // Over HTTP, resources support subscribe + server-push updates.
457 resources: resources_enabled.then_some(ResourcesCapability {
458 subscribe: true,
459 list_changed: false,
460 }),
461 },
462 server_info: ServerInfo {
463 name: "kberg".to_string(),
464 version: env!("CARGO_PKG_VERSION").to_string(),
465 },
466 })
467 .unwrap()
468 }
469
470 #[cfg(test)]
471 mod tests {
472 use super::*;
473 use crate::resource::{ResourceContents, ResourceDescriptor};
474 use crate::tool::{Tool, ToolCallResult, ToolKind};
475 use async_trait::async_trait;
476
477 struct Echo;
478 #[async_trait]
479 impl Tool for Echo {
480 fn name(&self) -> &'static str {
481 "echo"
482 }
483 fn description(&self) -> &'static str {
484 "read echo"
485 }
486 fn kind(&self) -> ToolKind {
487 ToolKind::Read
488 }
489 fn input_schema(&self) -> Value {
490 json!({ "type": "object", "properties": {} })
491 }
492 async fn call(&self, _args: Value) -> crate::Result<ToolCallResult> {
493 Ok(ToolCallResult::text("pong"))
494 }
495 }
496
497 fn registry() -> ToolRegistry {
498 let mut r = ToolRegistry::new();
499 r.register(Echo);
500 r
501 }
502
503 // id/params are consumed into the returned request (test helper).
504 #[allow(clippy::needless_pass_by_value)]
505 fn request(id: Option<Value>, method: &str, params: Value) -> JsonRpcRequest {
506 serde_json::from_value(json!({
507 "jsonrpc": "2.0",
508 "id": id,
509 "method": method,
510 "params": params,
511 }))
512 .unwrap()
513 }
514
515 #[tokio::test]
516 async fn dispatch_initialize_and_list_and_call() {
517 let r = registry();
518 let grants = HashSet::new();
519
520 let init = dispatch(
521 &r,
522 None,
523 SurfaceProjection::Full,
524 Some(&grants),
525 None,
526 request(Some(json!(1)), "initialize", Value::Null),
527 )
528 .await
529 .unwrap();
530 let v = serde_json::to_value(init).unwrap();
531 assert_eq!(v["result"]["serverInfo"]["name"], "kberg");
532 // No resource registry → no `resources` capability advertised.
533 assert!(v["result"]["capabilities"]["resources"].is_null());
534
535 let list = dispatch(
536 &r,
537 None,
538 SurfaceProjection::Full,
539 Some(&grants),
540 None,
541 request(Some(json!(2)), "tools/list", Value::Null),
542 )
543 .await
544 .unwrap();
545 let v = serde_json::to_value(list).unwrap();
546 assert_eq!(v["result"]["tools"][0]["name"], "echo");
547
548 let call = dispatch(
549 &r,
550 None,
551 SurfaceProjection::Full,
552 Some(&grants),
553 None,
554 request(
555 Some(json!(3)),
556 "tools/call",
557 json!({ "name": "echo", "arguments": {} }),
558 ),
559 )
560 .await
561 .unwrap();
562 let v = serde_json::to_value(call).unwrap();
563 assert_eq!(v["result"]["content"][0]["text"], "pong");
564 }
565
566 #[tokio::test]
567 async fn dispatch_notification_has_no_response() {
568 let r = registry();
569 let resp = dispatch(
570 &r,
571 None,
572 SurfaceProjection::Full,
573 None,
574 None,
575 request(None, "initialize", Value::Null),
576 )
577 .await;
578 assert!(resp.is_none());
579 }
580
581 #[tokio::test]
582 async fn dispatch_unknown_method_errors() {
583 let r = registry();
584 let resp = dispatch(
585 &r,
586 None,
587 SurfaceProjection::Full,
588 None,
589 None,
590 request(Some(json!(9)), "bogus", Value::Null),
591 )
592 .await
593 .unwrap();
594 let v = serde_json::to_value(resp).unwrap();
595 assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND);
596 }
597
598 fn resources() -> ResourceRegistry {
599 let res = ResourceRegistry::new();
600 res.register(
601 ResourceDescriptor::new("buffer://main.rs", "main.rs").with_mime_type("text/x-rust"),
602 |uri| async move { Ok(ResourceContents::text(uri, "fn main() {}")) },
603 );
604 res
605 }
606
607 #[tokio::test]
608 async fn dispatch_initialize_advertises_resources_when_configured() {
609 let r = registry();
610 let res = resources();
611 let init = dispatch(
612 &r,
613 Some(&res),
614 SurfaceProjection::Full,
615 None,
616 None,
617 request(Some(json!(1)), "initialize", Value::Null),
618 )
619 .await
620 .unwrap();
621 let v = serde_json::to_value(init).unwrap();
622 assert_eq!(v["result"]["capabilities"]["resources"]["subscribe"], true);
623 }
624
625 #[tokio::test]
626 async fn dispatch_resources_list_and_read() {
627 let r = registry();
628 let res = resources();
629
630 let list = dispatch(
631 &r,
632 Some(&res),
633 SurfaceProjection::Full,
634 None,
635 None,
636 request(Some(json!(2)), "resources/list", Value::Null),
637 )
638 .await
639 .unwrap();
640 let v = serde_json::to_value(list).unwrap();
641 assert_eq!(v["result"]["resources"][0]["uri"], "buffer://main.rs");
642 assert_eq!(v["result"]["resources"][0]["mimeType"], "text/x-rust");
643
644 let read = dispatch(
645 &r,
646 Some(&res),
647 SurfaceProjection::Full,
648 None,
649 None,
650 request(
651 Some(json!(3)),
652 "resources/read",
653 json!({ "uri": "buffer://main.rs" }),
654 ),
655 )
656 .await
657 .unwrap();
658 let v = serde_json::to_value(read).unwrap();
659 assert_eq!(v["result"]["contents"][0]["text"], "fn main() {}");
660 }
661
662 #[tokio::test]
663 async fn dispatch_resources_read_unknown_uri_errors() {
664 let r = registry();
665 let res = resources();
666 let read = dispatch(
667 &r,
668 Some(&res),
669 SurfaceProjection::Full,
670 None,
671 None,
672 request(
673 Some(json!(4)),
674 "resources/read",
675 json!({ "uri": "buffer://gone.rs" }),
676 ),
677 )
678 .await
679 .unwrap();
680 let v = serde_json::to_value(read).unwrap();
681 assert_eq!(v["error"]["code"], codes::RESOURCE_NOT_FOUND);
682 }
683
684 #[tokio::test]
685 async fn dispatch_resources_method_without_registry_is_method_not_found() {
686 let r = registry();
687 let resp = dispatch(
688 &r,
689 None,
690 SurfaceProjection::Full,
691 None,
692 None,
693 request(Some(json!(5)), "resources/list", Value::Null),
694 )
695 .await
696 .unwrap();
697 let v = serde_json::to_value(resp).unwrap();
698 assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND);
699 }
700 }
701