rust_web_server/mcp/mod.rs
1//! Model Context Protocol (MCP) server — HTTP Streamable HTTP transport.
2//!
3//! [`McpServer`] implements [`Application`] so it can be passed directly to
4//! [`Server::run`]. Unmatched requests fall through to the built-in [`App`]
5//! controller chain (static files, health probes, etc.).
6//!
7//! # Quick start
8//!
9//! ```rust,no_run
10//! use rust_web_server::server::Server;
11//! use rust_web_server::mcp::{McpServer, McpContent, PromptMessage};
12//! # fn main() {
13//! let mcp = McpServer::new("my-server", "1.0")
14//! // A tool: callable by the AI, like a function
15//! .tool(
16//! "echo",
17//! "Echo text back",
18//! r#"{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}"#,
19//! |args| {
20//! let text = rust_web_server::mcp::extract_arg(args, "text")
21//! .unwrap_or_else(|| "(nothing)".to_string());
22//! Ok(McpContent::text(text))
23//! },
24//! )
25//! // A resource: data the AI can read by URI
26//! .resource(
27//! "docs://{topic}",
28//! "Documentation",
29//! "Return documentation for a topic",
30//! |uri| Ok(McpContent::text(format!("Documentation for: {uri}"))),
31//! )
32//! // A prompt template: reusable message structures
33//! .prompt(
34//! "summarize",
35//! "Summarize the given text",
36//! |args| {
37//! let text = rust_web_server::mcp::extract_arg(args, "text")
38//! .unwrap_or_else(|| "some text".to_string());
39//! Ok(vec![PromptMessage::user(format!("Please summarize: {text}"))])
40//! },
41//! );
42//!
43//! // let (listener, pool) = Server::setup().unwrap();
44//! // Server::run(listener, pool, mcp);
45//! # }
46//! ```
47//!
48//! # MCP endpoint
49//!
50//! All JSON-RPC messages are sent as `POST /mcp` (override with [`.at()`](McpServer::at)).
51//! The server implements the [MCP 2024-11-05 specification](https://spec.modelcontextprotocol.io).
52//!
53//! `GET /mcp` opens a Server-Sent Events stream for server → client push —
54//! see [`McpServer::notify`] and the module docs' SSE section below.
55//!
56//! # SSE streaming transport
57//!
58//! A client that sends `GET /mcp` (instead of `POST`) gets back a
59//! `text/event-stream` response that stays open indefinitely. Call
60//! [`McpServer::notify`] from anywhere (a background thread, another request's
61//! handler, ...) to push a JSON-RPC notification to every currently-connected
62//! SSE client:
63//!
64//! ```rust,no_run
65//! use rust_web_server::mcp::McpServer;
66//!
67//! let server = McpServer::new("my-server", "1.0");
68//! // Elsewhere, e.g. after some background job finishes:
69//! server.notify("notifications/message", Some(r#"{"level":"info","data":"job done"}"#));
70//! ```
71//!
72//! Idle connections receive a `: keep-alive` SSE comment every 15 seconds so
73//! intermediate proxies don't time them out; this doubles as the mechanism
74//! that detects a client has disconnected (the next write attempt fails and
75//! the connection is dropped). A client whose event buffer fills up (32
76//! pending frames, unconsumed) is treated the same as a disconnected one and
77//! dropped from the broadcast list — [`McpServer::notify`] never blocks the
78//! calling thread waiting on a slow reader.
79//!
80//! This transport is only wired up for the plain HTTP/1.1 path
81//! (`Server::run`/`Server::process`) — same scope as `Response::stream_pipe`
82//! generally, which the HTTP/2 (`h2_handler`) and HTTP/3 (`h3_handler`)
83//! handlers don't yet support for *any* response, not just this one.
84//!
85//! # Environment variables
86//!
87//! None — configure the server programmatically via the builder.
88
89mod json_rpc;
90
91#[cfg(test)]
92mod tests;
93
94use std::collections::HashMap;
95#[cfg(feature = "http2")]
96use std::future::Future;
97#[cfg(feature = "http2")]
98use std::pin::Pin;
99use std::sync::mpsc::{self, SyncSender};
100use std::sync::{Arc, Mutex, RwLock};
101use std::time::Duration;
102
103use crate::app::App;
104use crate::application::Application;
105use crate::core::New;
106use crate::header::Header;
107use crate::mime_type::MimeType;
108use crate::range::Range;
109use crate::request::Request;
110use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
111use crate::server::ConnectionInfo;
112
113const PROTOCOL_VERSION: &str = "2024-11-05";
114
115// ── public content types ──────────────────────────────────────────────────────
116
117/// Content returned by tool and resource handlers.
118///
119/// Create with [`McpContent::text`] (plain text or JSON strings),
120/// [`McpContent::json`] (marks MIME type as `application/json`),
121/// [`McpContent::image`] (base64-encoded binary image data), or
122/// [`McpContent::embedded`] (a resource embedded inline in a tool response).
123#[derive(Clone, Debug)]
124pub struct McpContent {
125 /// `"text"`, `"image"`, or `"resource"`.
126 pub kind: &'static str,
127 /// The content string — text for `"text"`, base64 data for `"image"`,
128 /// or the embedded resource's text for `"resource"`.
129 pub text: String,
130 /// Optional MIME type override (default `"text/plain"` for `"text"`;
131 /// required in practice for `"image"`/`"resource"`, set by their
132 /// constructors).
133 pub mime_type: Option<String>,
134 /// The resource URI — only set (and only serialized) for `"resource"`.
135 pub uri: Option<String>,
136}
137
138impl McpContent {
139 /// Plain-text content.
140 pub fn text(s: impl Into<String>) -> Self {
141 McpContent { kind: "text", text: s.into(), mime_type: None, uri: None }
142 }
143
144 /// JSON content — sets `mimeType` to `application/json`.
145 pub fn json(s: impl Into<String>) -> Self {
146 McpContent { kind: "text", text: s.into(), mime_type: Some("application/json".to_string()), uri: None }
147 }
148
149 /// Image content (screenshot, chart, generated art) — `data` is base64-encoded
150 /// binary and `mime_type` is e.g. `"image/png"`.
151 pub fn image(data: impl Into<String>, mime_type: impl Into<String>) -> Self {
152 McpContent { kind: "image", text: data.into(), mime_type: Some(mime_type.into()), uri: None }
153 }
154
155 /// A resource embedded inline in a tool response, as opposed to one a
156 /// client fetches separately via `resources/read`.
157 pub fn embedded(uri: impl Into<String>, text: impl Into<String>, mime_type: impl Into<String>) -> Self {
158 McpContent { kind: "resource", text: text.into(), mime_type: Some(mime_type.into()), uri: Some(uri.into()) }
159 }
160
161 fn to_content_json(&self) -> String {
162 match self.kind {
163 "image" => format!(
164 r#"{{"type":"image","data":"{}","mimeType":"{}"}}"#,
165 json_escape(&self.text),
166 json_escape(self.mime_type.as_deref().unwrap_or("application/octet-stream")),
167 ),
168 "resource" => format!(
169 r#"{{"type":"resource","resource":{{"uri":"{}","mimeType":"{}","text":"{}"}}}}"#,
170 json_escape(self.uri.as_deref().unwrap_or("")),
171 json_escape(self.mime_type.as_deref().unwrap_or("text/plain")),
172 json_escape(&self.text),
173 ),
174 _ => format!(r#"{{"type":"text","text":"{}"}}"#, json_escape(&self.text)),
175 }
176 }
177
178 fn mime(&self) -> &str {
179 self.mime_type.as_deref().unwrap_or("text/plain")
180 }
181}
182
183/// A single message in a prompt response.
184#[derive(Clone, Debug)]
185pub struct PromptMessage {
186 /// `"user"` or `"assistant"`.
187 pub role: &'static str,
188 /// The message content.
189 pub content: McpContent,
190}
191
192impl PromptMessage {
193 /// Build a user-role message.
194 pub fn user(text: impl Into<String>) -> Self {
195 PromptMessage { role: "user", content: McpContent::text(text) }
196 }
197
198 /// Build an assistant-role message.
199 pub fn assistant(text: impl Into<String>) -> Self {
200 PromptMessage { role: "assistant", content: McpContent::text(text) }
201 }
202
203 fn to_json(&self) -> String {
204 format!(
205 r#"{{"role":"{}","content":{}}}"#,
206 self.role,
207 self.content.to_content_json(),
208 )
209 }
210}
211
212/// Argument definition for a prompt template.
213#[derive(Clone)]
214pub struct PromptArgDef {
215 pub name: String,
216 pub description: String,
217 pub required: bool,
218}
219
220impl PromptArgDef {
221 pub fn required(name: impl Into<String>, description: impl Into<String>) -> Self {
222 PromptArgDef { name: name.into(), description: description.into(), required: true }
223 }
224
225 pub fn optional(name: impl Into<String>, description: impl Into<String>) -> Self {
226 PromptArgDef { name: name.into(), description: description.into(), required: false }
227 }
228}
229
230// ── McpContext ────────────────────────────────────────────────────────────────
231
232/// Per-request context passed to tool handlers registered via
233/// [`McpServer::tool_with_context`] — caller identity and session info that a
234/// plain `Fn(&str) -> ...` tool handler has no way to see.
235///
236/// Constructed in [`McpServer::execute`] from the current request's headers
237/// plus whatever `clientInfo` was recorded for this session at `initialize`
238/// time (see [`McpServer::handle_request_with_context`]).
239#[derive(Debug, Clone, Default)]
240pub struct McpContext {
241 /// `clientInfo.name` sent in this session's `initialize` call, if the
242 /// client sent one and this request carries a recognized `Mcp-Session-Id`.
243 pub client_name: Option<String>,
244 /// `clientInfo.version` sent in this session's `initialize` call, under
245 /// the same conditions as `client_name`.
246 pub client_version: Option<String>,
247 /// The `Mcp-Session-Id` header on this request, if present — the value
248 /// the server minted and returned in the `initialize` response header
249 /// for this session (see the module docs' Sessions section).
250 pub session_id: Option<String>,
251 /// Verified OAuth 2.0 / OIDC JWT claims as a JSON string (serialized
252 /// [`crate::sso::OidcClaims`]), populated when [`McpServer::require_oauth`]
253 /// (`sso` feature) is configured and this request's bearer token
254 /// verified successfully. `None` if OAuth isn't configured, or for a
255 /// hand-built context (e.g. `handle_request` in tests) that bypassed
256 /// `McpServer::execute`'s auth check entirely.
257 pub auth_claims: Option<String>,
258 /// The raw JSON value of `params._meta.progressToken` from the
259 /// triggering `tools/call` request, if the client sent one — a spec
260 /// `string | number`, so this is stored pre-rendered (already correctly
261 /// quoted if it's a string) rather than decoded, and spliced back
262 /// verbatim by [`Self::report_progress`]. `None` for anything other than
263 /// a `tools/call` whose caller asked for progress updates.
264 pub progress_token: Option<String>,
265 /// Shared handle back to the owning [`McpServer`]'s SSE broadcast list,
266 /// used by [`Self::report_progress`]. Not `pub` — this is plumbing, not
267 /// part of the context data a handler reads. `None` for a context built
268 /// by hand (e.g. via [`McpServer::handle_request_with_context`] in a
269 /// test) rather than through [`McpServer::execute`], in which case
270 /// `report_progress` silently no-ops — there's no live server to
271 /// broadcast through.
272 sse_clients: Option<Arc<Mutex<Vec<SseClient>>>>,
273 /// Shared flag flipped by a `notifications/cancelled` referencing this
274 /// `tools/call`'s request id, checked by [`Self::is_cancelled`]. Not
275 /// `pub` — see that method. `None` for anything other than a live
276 /// `tools/call` dispatched through [`McpServer::execute`]/`handle_request_with_context`.
277 cancellation: Option<Arc<std::sync::atomic::AtomicBool>>,
278 /// Whether this session's `initialize` call declared
279 /// `params.capabilities.sampling`, checked by [`Self::sample`] before
280 /// sending a request the client never said it could answer.
281 sampling_supported: bool,
282 /// Shared handle to the owning [`McpServer`]'s in-flight
283 /// server-initiated-request registry (`sampling/createMessage`,
284 /// `roots/list`), used by [`Self::sample`]/[`Self::list_roots`]. Not
285 /// `pub` — see `sse_clients` for the same "plumbing, not context data"
286 /// reasoning; `None` under the same conditions.
287 pending_replies: Option<Arc<Mutex<HashMap<String, mpsc::Sender<Result<String, String>>>>>>,
288 /// Whether this session's `initialize` call declared
289 /// `params.capabilities.roots`, checked by [`Self::list_roots`] before
290 /// sending a request the client never said it could answer.
291 roots_supported: bool,
292 /// Shared handle to the owning [`McpServer`]'s session registry, used
293 /// by [`Self::list_roots`] to read/write this session's cached roots
294 /// list. Not `pub` — plumbing, not context data.
295 sessions: Option<Arc<Mutex<HashMap<String, StoredClientInfo>>>>,
296}
297
298impl McpContext {
299 /// Push a `notifications/progress` event over the SSE channel for this
300 /// request's `progressToken`, if the client asked for progress updates
301 /// (`params._meta.progressToken` on the triggering `tools/call`) and
302 /// this context was built through a live [`McpServer`] (via `execute()`,
303 /// not a bare `McpContext { .. }` — see the `sse_clients` field doc).
304 ///
305 /// Silently does nothing in either case — a handler doesn't need to
306 /// branch on whether progress reporting is actually wired up before
307 /// calling this; it's always safe to call.
308 ///
309 /// `total` and `message` are both optional, matching the spec's
310 /// `notifications/progress` shape: `{"progressToken":...,"progress":...,
311 /// "total":...,"message":"..."}` (with `total`/`message` omitted when not
312 /// given here).
313 ///
314 /// ```rust,no_run
315 /// use rust_web_server::mcp::{McpContent, McpServer};
316 ///
317 /// let server = McpServer::new("my-server", "1.0")
318 /// .tool_with_context("long_job", "Do something slow", "{}", |ctx, _args| {
319 /// ctx.report_progress(0.0, Some(100.0), Some("starting"));
320 /// // ... do work ...
321 /// ctx.report_progress(100.0, Some(100.0), Some("done"));
322 /// Ok(McpContent::text("done"))
323 /// });
324 /// ```
325 pub fn report_progress(&self, progress: f64, total: Option<f64>, message: Option<&str>) {
326 let (Some(token), Some(sse_clients)) = (&self.progress_token, &self.sse_clients) else {
327 return;
328 };
329
330 let total_field = match total {
331 Some(t) => format!(r#","total":{t}"#),
332 None => String::new(),
333 };
334 let message_field = match message {
335 Some(m) => format!(r#","message":"{}""#, json_escape(m)),
336 None => String::new(),
337 };
338 let params = format!(
339 r#"{{"progressToken":{token},"progress":{progress}{total_field}{message_field}}}"#
340 );
341 let json = render_notification("notifications/progress", Some(¶ms));
342 broadcast_sse_to(sse_clients, &json);
343 }
344
345 /// `true` if the client sent `notifications/cancelled` referencing this
346 /// `tools/call`'s request id.
347 ///
348 /// Rust cannot forcibly interrupt a running synchronous closure — same
349 /// limitation [`crate::timeout::with_timeout`] documents — so this is
350 /// *cooperative* cancellation: nothing happens automatically. A handler
351 /// that does its own iterative work (processing a batch of items, say)
352 /// can check this between steps and return early; a handler that never
353 /// checks it runs to completion regardless, exactly as before this
354 /// existed.
355 ///
356 /// Always `false` if the client never sent a cancellation, if this
357 /// wasn't dispatched as a `tools/call` (the only method cancellation
358 /// applies to), or if `ctx` was built without a live server behind it.
359 ///
360 /// ```rust,no_run
361 /// use rust_web_server::mcp::{McpContent, McpServer};
362 ///
363 /// let server = McpServer::new("my-server", "1.0")
364 /// .tool_with_context("process_batch", "Process many items", "{}", |ctx, _args| {
365 /// for i in 0..1_000_000 {
366 /// if ctx.is_cancelled() {
367 /// return Err("cancelled by client".to_string());
368 /// }
369 /// // ... process item i ...
370 /// }
371 /// Ok(McpContent::text("done"))
372 /// });
373 /// ```
374 pub fn is_cancelled(&self) -> bool {
375 self.cancellation.as_ref().is_some_and(|c| c.load(std::sync::atomic::Ordering::Relaxed))
376 }
377
378 /// Ask the connected client to run LLM inference ("sampling") and block
379 /// until it replies or `timeout` elapses — the server-initiated half of
380 /// `sampling/createMessage`. This reverses the usual request direction:
381 /// the server sends a JSON-RPC *request* to the client over the `GET
382 /// /mcp` SSE stream, and the client answers with its own `POST /mcp`
383 /// carrying a JSON-RPC *response* (no `method`) that
384 /// [`McpServer::handle_request_with_context`] recognizes and routes
385 /// back here instead of treating it as an invalid request.
386 ///
387 /// Blocking (not `async fn`) is deliberate: tool handlers in this crate
388 /// are plain synchronous closures — there is no async tool handler
389 /// support yet (MCP_TODO.md's TODO-17) for this to `.await` inside of.
390 /// The calling thread parks on the response channel for up to
391 /// `timeout`; on a thread-pool server this ties up one worker for that
392 /// long, same tradeoff `crate::timeout::with_timeout` already accepts
393 /// for bounding a slow handler's *caller*.
394 ///
395 /// Fails fast, before sending anything, if: the client's `initialize`
396 /// call never declared `capabilities.sampling` (it wouldn't know what
397 /// to do with the request); this request has no session id (nothing to
398 /// address the request to); or `ctx` has no live server behind it.
399 /// Otherwise fails with a timeout error if the client never responds —
400 /// including if it simply has no `GET /mcp` SSE connection open for
401 /// this session, since there's no separate "not connected" signal.
402 ///
403 /// ```rust,no_run
404 /// use rust_web_server::mcp::{McpServer, PromptMessage, SamplingRequest};
405 /// use std::time::Duration;
406 ///
407 /// let server = McpServer::new("my-server", "1.0")
408 /// .tool_with_context("ask_llm", "Ask the connected client's model a question", "{}", |ctx, _args| {
409 /// let response = ctx.sample(
410 /// SamplingRequest {
411 /// messages: vec![PromptMessage::user("What is 2+2?")],
412 /// max_tokens: 100,
413 /// system_prompt: None,
414 /// },
415 /// Duration::from_secs(30),
416 /// )?;
417 /// Ok(response.content)
418 /// });
419 /// ```
420 pub fn sample(&self, request: SamplingRequest, timeout: Duration) -> Result<SamplingResponse, String> {
421 if !self.sampling_supported {
422 return Err("the connected client did not declare sampling support in its initialize capabilities".to_string());
423 }
424
425 let messages_json: Vec<String> = request.messages.iter().map(|m| m.to_json()).collect();
426 let system_prompt_field = match &request.system_prompt {
427 Some(s) => format!(r#","systemPrompt":"{}""#, json_escape(s)),
428 None => String::new(),
429 };
430 let params = format!(
431 r#"{{"messages":[{}],"maxTokens":{}{system_prompt_field}}}"#,
432 messages_json.join(","), request.max_tokens,
433 );
434
435 let result_json = self.send_and_wait("sampling/createMessage", Some(¶ms), timeout)?;
436 parse_sampling_response(&result_json)
437 }
438
439 /// Ask the connected client which filesystem roots (workspace
440 /// directories, mounted volumes, ...) it has access to — useful for a
441 /// file-system-aware tool that should only operate within the client's
442 /// declared workspace rather than the whole filesystem.
443 ///
444 /// Like [`Self::sample`], this reverses the usual request direction (a
445 /// server-initiated `roots/list` request over SSE, answered by the
446 /// client's own `POST /mcp`) and blocks the calling thread for up to
447 /// `timeout` — see `sample`'s docs for why blocking is deliberate here.
448 ///
449 /// The result is cached per session: the first call after `initialize`
450 /// (or after the client sends `notifications/roots/list_changed`, which
451 /// invalidates the cache) does a live round trip; later calls in the
452 /// same session return the cached list without sending anything. A tool
453 /// handler can call this on every invocation without worrying about
454 /// spamming the client with redundant `roots/list` requests.
455 ///
456 /// Fails fast, before sending anything, under the same conditions as
457 /// [`Self::sample`] (translated to roots): the client never declared
458 /// `capabilities.roots`, this request has no session id, or `ctx` has
459 /// no live server behind it.
460 ///
461 /// ```rust,no_run
462 /// use rust_web_server::mcp::McpServer;
463 /// use std::time::Duration;
464 ///
465 /// let server = McpServer::new("my-server", "1.0")
466 /// .tool_with_context("list_workspace_files", "List files in the client's workspace", "{}", |ctx, _args| {
467 /// let roots = ctx.list_roots(Duration::from_secs(10))?;
468 /// let names: Vec<String> = roots.iter().map(|r| r.uri.clone()).collect();
469 /// Ok(rust_web_server::mcp::McpContent::text(names.join(", ")))
470 /// });
471 /// ```
472 pub fn list_roots(&self, timeout: Duration) -> Result<Vec<McpRoot>, String> {
473 if !self.roots_supported {
474 return Err("the connected client did not declare roots support in its initialize capabilities".to_string());
475 }
476
477 if let (Some(session_id), Some(sessions)) = (&self.session_id, &self.sessions) {
478 if let Some(cached) = sessions.lock().unwrap().get(session_id).and_then(|info| info.roots.clone()) {
479 return Ok(cached);
480 }
481 }
482
483 let result_json = self.send_and_wait("roots/list", None, timeout)?;
484 let roots = parse_roots_response(&result_json)?;
485
486 if let (Some(session_id), Some(sessions)) = (&self.session_id, &self.sessions) {
487 if let Some(info) = sessions.lock().unwrap().get_mut(session_id) {
488 info.roots = Some(roots.clone());
489 }
490 }
491
492 Ok(roots)
493 }
494
495 /// Send a JSON-RPC request to the client for this context's session and
496 /// block for the reply — the shared send-and-correlate mechanic behind
497 /// [`Self::sample`] and [`Self::list_roots`]. Returns the raw `result`
498 /// JSON string (not yet parsed into either method's own response type)
499 /// on success.
500 fn send_and_wait(&self, method: &str, params_json: Option<&str>, timeout: Duration) -> Result<String, String> {
501 let session_id = self.session_id.clone()
502 .ok_or_else(|| format!("{method} requires a session (Mcp-Session-Id)"))?;
503 let sse_clients = self.sse_clients.as_ref()
504 .ok_or_else(|| format!("{method} requires a live server (this context has none)"))?;
505 let pending = self.pending_replies.as_ref()
506 .ok_or_else(|| format!("{method} requires a live server (this context has none)"))?;
507
508 let request_id = format!("\"{}\"", crate::request_id::generate_request_id());
509 let (tx, rx) = mpsc::channel::<Result<String, String>>();
510 pending.lock().unwrap().insert(request_id.clone(), tx);
511
512 let params_field = match params_json {
513 Some(p) => format!(r#","params":{p}"#),
514 None => String::new(),
515 };
516 let rpc = format!(r#"{{"jsonrpc":"2.0","id":{request_id},"method":"{method}"{params_field}}}"#);
517
518 send_sse_to_sessions(sse_clients, &[session_id], &rpc);
519
520 let outcome = rx.recv_timeout(timeout);
521 pending.lock().unwrap().remove(&request_id);
522
523 match outcome {
524 Ok(Ok(result_json)) => Ok(result_json),
525 Ok(Err(e)) => Err(e),
526 Err(_) => Err(format!("{method} timed out waiting for the client's response")),
527 }
528 }
529}
530
531/// A request to ask the connected MCP client to run LLM inference
532/// ("sampling") — build one and pass it to [`McpContext::sample`].
533///
534/// `messages` reuses [`PromptMessage`] rather than introducing a near-identical
535/// type — the spec's sampling message shape (`{"role":...,"content":{"type":"text",...}}`)
536/// is exactly what `PromptMessage` already models.
537pub struct SamplingRequest {
538 pub messages: Vec<PromptMessage>,
539 pub max_tokens: u32,
540 /// An optional system prompt to steer the client's model, per the
541 /// spec's `params.systemPrompt`.
542 pub system_prompt: Option<String>,
543}
544
545/// The client's answer to a [`SamplingRequest`], returned by
546/// [`McpContext::sample`].
547pub struct SamplingResponse {
548 /// Almost always `"assistant"`, per spec.
549 pub role: String,
550 pub content: McpContent,
551 /// The model the client actually used to generate this response, e.g.
552 /// `"claude-3-sonnet-20240307"`.
553 pub model: String,
554 /// Why the model stopped, e.g. `"endTurn"`, `"maxTokens"`, `"stopSequence"`.
555 /// `None` if the client didn't report one.
556 pub stop_reason: Option<String>,
557}
558
559/// Parse a `sampling/createMessage` response's `result` object (already
560/// extracted from the enclosing JSON-RPC envelope) into a [`SamplingResponse`].
561fn parse_sampling_response(result_json: &str) -> Result<SamplingResponse, String> {
562 let role = json_rpc::extract_str(result_json, "role").unwrap_or_else(|| "assistant".to_string());
563 let content_json = json_rpc::extract_raw(result_json, "content")
564 .ok_or_else(|| "sampling response missing content".to_string())?;
565 let kind = json_rpc::extract_str(&content_json, "type").unwrap_or_else(|| "text".to_string());
566 let content = match kind.as_str() {
567 "image" => McpContent::image(
568 json_rpc::extract_str(&content_json, "data").unwrap_or_default(),
569 json_rpc::extract_str(&content_json, "mimeType").unwrap_or_default(),
570 ),
571 _ => McpContent::text(json_rpc::extract_str(&content_json, "text").unwrap_or_default()),
572 };
573 let model = json_rpc::extract_str(result_json, "model").unwrap_or_default();
574 let stop_reason = json_rpc::extract_str(result_json, "stopReason");
575 Ok(SamplingResponse { role, content, model, stop_reason })
576}
577
578/// One filesystem root a client has access to, returned by
579/// [`McpContext::list_roots`] — a workspace directory, a mounted volume,
580/// or similar. `uri` is typically a `file://` URI per spec.
581#[derive(Debug, Clone)]
582pub struct McpRoot {
583 pub uri: String,
584 /// A human-readable label for the root, if the client provided one.
585 pub name: Option<String>,
586}
587
588/// Parse a `roots/list` response's `result` object (already extracted from
589/// the enclosing JSON-RPC envelope) into a `Vec<McpRoot>`.
590fn parse_roots_response(result_json: &str) -> Result<Vec<McpRoot>, String> {
591 let roots_array = json_rpc::extract_raw(result_json, "roots")
592 .ok_or_else(|| "roots/list response missing roots".to_string())?;
593 let roots = json_rpc::split_array_elements(&roots_array).iter().map(|elem| {
594 McpRoot {
595 uri: json_rpc::extract_str(elem, "uri").unwrap_or_default(),
596 name: json_rpc::extract_str(elem, "name"),
597 }
598 }).collect();
599 Ok(roots)
600}
601
602/// Configuration installed by [`McpServer::require_oauth`] (`sso` feature).
603/// Reuses [`crate::sso::jwks::JwksCache`] and [`crate::sso::OidcProvider`]
604/// directly rather than re-implementing JWKS fetch/verify — see
605/// `require_oauth`'s doc comment for the full rationale.
606#[cfg(feature = "sso")]
607struct OAuthConfig {
608 jwks: crate::sso::jwks::JwksCache,
609 provider: crate::sso::OidcProvider,
610 audience: String,
611}
612
613/// `clientInfo` recorded for one session at `initialize` time, looked up by
614/// `Mcp-Session-Id` for later requests in the same session. See
615/// `McpServer`'s `sessions` field doc comment for the unbounded-growth caveat.
616#[derive(Debug, Clone, Default)]
617struct StoredClientInfo {
618 name: Option<String>,
619 version: Option<String>,
620 /// Whether this session's `initialize` call declared `params.capabilities.sampling`
621 /// (spec: the client, not the server, declares sampling support — the
622 /// presence of the key is what matters, not its contents, which are
623 /// normally an empty object). Checked by [`McpContext::sample`] before
624 /// sending a `sampling/createMessage` request that the client would
625 /// have no idea how to answer.
626 supports_sampling: bool,
627 /// Whether this session's `initialize` call declared `params.capabilities.roots`,
628 /// checked by [`McpContext::list_roots`] the same way `supports_sampling` gates `sample`.
629 supports_roots: bool,
630 /// Cached result of this session's last `roots/list` round trip, read
631 /// and written by [`McpContext::list_roots`]. `None` means "never
632 /// fetched, or invalidated" — set back to `None` when this session
633 /// sends `notifications/roots/list_changed`, so the next `list_roots`
634 /// call does a fresh round trip instead of serving stale data.
635 roots: Option<Vec<McpRoot>>,
636}
637
638// ── ToolAnnotations ───────────────────────────────────────────────────────────
639
640/// Behavioral hints for a tool, per the MCP 2025-03-26 spec's tool
641/// annotations. Clients (Claude Desktop and others) use these to decide
642/// whether to warn or ask for confirmation before calling a tool — e.g. skip
643/// confirmation for a read-only tool, or warn before a destructive one.
644///
645/// Every field is a *hint*, not something this server enforces or verifies —
646/// nothing stops a handler registered with `read_only_hint: Some(true)` from
647/// actually writing to disk. A well-behaved server sets these accurately;
648/// a client is free to ignore them or ask for confirmation anyway.
649///
650/// Register with [`McpServer::tool_annotated`]. Build one with plain struct
651/// syntax — every field defaults to `None` (no hint given, the client's own
652/// default applies):
653///
654/// ```rust
655/// use rust_web_server::mcp::ToolAnnotations;
656///
657/// let destructive = ToolAnnotations {
658/// destructive_hint: Some(true),
659/// read_only_hint: Some(false),
660/// ..Default::default()
661/// };
662/// ```
663#[derive(Debug, Clone, Copy, Default)]
664pub struct ToolAnnotations {
665 /// The tool does not modify its environment.
666 pub read_only_hint: Option<bool>,
667 /// The tool may perform destructive updates (only meaningful when
668 /// `read_only_hint` is not `Some(true)`).
669 pub destructive_hint: Option<bool>,
670 /// Calling the tool repeatedly with the same arguments has no additional
671 /// effect beyond the first call.
672 pub idempotent_hint: Option<bool>,
673 /// The tool may interact with an open-ended set of external entities
674 /// (e.g. web search), as opposed to a fixed, closed set.
675 pub open_world_hint: Option<bool>,
676}
677
678impl ToolAnnotations {
679 /// Render as a JSON object containing only the hints that are `Some`,
680 /// using the spec's camelCase key names. Returns `"{}"` if every field
681 /// is `None`.
682 fn to_json(self) -> String {
683 let mut fields = Vec::with_capacity(4);
684 if let Some(v) = self.read_only_hint {
685 fields.push(format!(r#""readOnlyHint":{v}"#));
686 }
687 if let Some(v) = self.destructive_hint {
688 fields.push(format!(r#""destructiveHint":{v}"#));
689 }
690 if let Some(v) = self.idempotent_hint {
691 fields.push(format!(r#""idempotentHint":{v}"#));
692 }
693 if let Some(v) = self.open_world_hint {
694 fields.push(format!(r#""openWorldHint":{v}"#));
695 }
696 format!("{{{}}}", fields.join(","))
697 }
698}
699
700// ── LogLevel ──────────────────────────────────────────────────────────────────
701
702/// RFC 5424 syslog severity levels, as used by the MCP `logging/setLevel`
703/// request and `notifications/message` log entries — ordered from most to
704/// least verbose so `level < min_level` comparisons work directly (this
705/// relies on declaration order matching severity order; don't reorder the
706/// variants).
707///
708/// ```rust
709/// use rust_web_server::mcp::LogLevel;
710///
711/// assert!(LogLevel::Debug < LogLevel::Warning);
712/// assert!(LogLevel::Emergency > LogLevel::Error);
713/// assert_eq!(LogLevel::parse("warning"), Some(LogLevel::Warning));
714/// assert_eq!(LogLevel::Warning.as_str(), "warning");
715/// ```
716#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
717pub enum LogLevel {
718 Debug,
719 Info,
720 Notice,
721 Warning,
722 Error,
723 Critical,
724 Alert,
725 Emergency,
726}
727
728impl LogLevel {
729 /// Parse the MCP spec's lowercase level name. Returns `None` for
730 /// anything that isn't one of the eight recognized levels.
731 pub fn parse(s: &str) -> Option<Self> {
732 match s {
733 "debug" => Some(LogLevel::Debug),
734 "info" => Some(LogLevel::Info),
735 "notice" => Some(LogLevel::Notice),
736 "warning" => Some(LogLevel::Warning),
737 "error" => Some(LogLevel::Error),
738 "critical" => Some(LogLevel::Critical),
739 "alert" => Some(LogLevel::Alert),
740 "emergency" => Some(LogLevel::Emergency),
741 _ => None,
742 }
743 }
744
745 /// The MCP spec's lowercase level name, e.g. `"warning"`.
746 pub fn as_str(self) -> &'static str {
747 match self {
748 LogLevel::Debug => "debug",
749 LogLevel::Info => "info",
750 LogLevel::Notice => "notice",
751 LogLevel::Warning => "warning",
752 LogLevel::Error => "error",
753 LogLevel::Critical => "critical",
754 LogLevel::Alert => "alert",
755 LogLevel::Emergency => "emergency",
756 }
757 }
758}
759
760// ── internal handler registrations ───────────────────────────────────────────
761
762type ToolFn = Arc<dyn Fn(McpContext, &str) -> Result<McpContent, String> + Send + Sync>;
763type ResourceFn = Arc<dyn Fn(&str) -> Result<McpContent, String> + Send + Sync>;
764type PromptFn = Arc<dyn Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync>;
765/// `Fn(argument_name, partial_value) -> candidate completion strings`.
766type CompletionFn = Arc<dyn Fn(&str, &str) -> Result<Vec<String>, String> + Send + Sync>;
767/// An async tool handler, boxed so `AsyncToolDef` can store handlers with
768/// different concrete `Future` types uniformly. The `Future` itself doesn't
769/// need `Send` for [`crate::async_bridge::block_on_isolated`] (which drives
770/// it to completion without ever moving it across a thread boundary), but
771/// it's required here anyway for consistency with [`crate::async_state`]'s
772/// equivalent handler type — the common case in practice, and one less
773/// surprise for anyone switching between the two async handler styles.
774#[cfg(feature = "http2")]
775type AsyncToolFn = Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<McpContent, String>> + Send>> + Send + Sync>;
776
777#[derive(Clone)]
778struct ToolDef {
779 name: String,
780 description: String,
781 input_schema: String,
782 annotations: Option<ToolAnnotations>,
783 handler: ToolFn,
784}
785
786/// A tool registered via [`McpServer::async_tool`]/[`McpServer::register_async_tool`]
787/// — stored separately from [`ToolDef`] rather than unifying the two behind
788/// a handler enum, since sync and async tools are otherwise identical and
789/// keeping them apart touches far less of the existing (already
790/// substantial) tool-handling code. `tools/list` and `tools/call` merge
791/// both collections transparently; from the client's perspective there is
792/// no distinction.
793#[cfg(feature = "http2")]
794#[derive(Clone)]
795struct AsyncToolDef {
796 name: String,
797 description: String,
798 input_schema: String,
799 annotations: Option<ToolAnnotations>,
800 handler: AsyncToolFn,
801}
802
803#[derive(Clone)]
804struct ResourceDef {
805 uri_template: String,
806 name: String,
807 description: String,
808 handler: ResourceFn,
809}
810
811#[derive(Clone)]
812struct PromptDef {
813 name: String,
814 description: String,
815 arguments: Vec<PromptArgDef>,
816 handler: PromptFn,
817}
818
819/// One `.completion()` registration — completion candidates for a single
820/// named argument of a single tool or prompt. `ref_type` is the short form
821/// passed to `.completion()` (e.g. `"tool"`, `"prompt"`), matched against the
822/// request's `ref.type` (e.g. `"ref/tool"`) with the `"ref/"` prefix
823/// stripped, not the raw wire value.
824#[derive(Clone)]
825struct CompletionDef {
826 ref_type: String,
827 ref_name: String,
828 handler: CompletionFn,
829}
830
831// ── McpServer ─────────────────────────────────────────────────────────────────
832
833/// An HTTP server that implements the MCP 2024-11-05 protocol.
834///
835/// Register tools, resources, and prompts with the builder methods, then pass
836/// the server to [`Server::run`] (or [`Server::run_tls`]) as an [`Application`].
837/// Requests that do not match the MCP endpoint fall through to the built-in
838/// [`App`] controller chain.
839#[derive(Clone)]
840pub struct McpServer {
841 server_name: String,
842 server_version: String,
843 path: String,
844 /// `Arc<RwLock<_>>` (not a plain `Vec`) so a running server's tool list
845 /// can be mutated at runtime — see [`Self::register_tool`]/[`Self::remove_tool`]
846 /// — and every clone of this `McpServer` (each connection thread gets
847 /// one) sees the same live list.
848 tools: Arc<RwLock<Vec<ToolDef>>>,
849 /// Tools registered via [`Self::async_tool`]/[`Self::register_async_tool`],
850 /// bridged into synchronous [`Application::execute`] via
851 /// [`crate::async_bridge::block_on_isolated`] at call time. Kept
852 /// separate from `tools` rather than merged into one collection — see
853 /// [`AsyncToolDef`]'s doc comment.
854 #[cfg(feature = "http2")]
855 async_tools: Arc<RwLock<Vec<AsyncToolDef>>>,
856 resources: Arc<RwLock<Vec<ResourceDef>>>,
857 prompts: Arc<RwLock<Vec<PromptDef>>>,
858 /// Argument completion providers registered via [`Self::completion`].
859 /// `initialize` advertises the `completions` capability iff this is
860 /// non-empty at that moment.
861 completions: Arc<RwLock<Vec<CompletionDef>>>,
862 fallback: Option<Arc<dyn Application + Send + Sync>>,
863 auth_token: Option<String>,
864 /// Installed by [`Self::require_oauth`] (`sso` feature) — an alternative
865 /// to `auth_token` that verifies a client-supplied OAuth 2.0 / OIDC
866 /// bearer JWT via JWKS instead of comparing against one static shared
867 /// secret. If both are somehow configured, OAuth verification takes
868 /// precedence and `auth_token` is ignored (see [`Self::require_oauth`]).
869 #[cfg(feature = "sso")]
870 oauth: Option<Arc<OAuthConfig>>,
871 /// Max items per page for `tools/list`/`resources/list`/`prompts/list`,
872 /// set via [`Self::page_size`]. `None` (the default) means no pagination
873 /// — every item comes back in one response, same as before pagination
874 /// existed.
875 page_size: Option<usize>,
876 /// `clientInfo` recorded per `Mcp-Session-Id`, minted at `initialize` time.
877 /// `Arc<Mutex<_>>` so every clone of this `McpServer` shares one map.
878 ///
879 /// This map only grows — nothing ever removes an entry, since there's no
880 /// session-termination signal in the MCP Streamable HTTP transport to key
881 /// eviction off of. Fine for the expected usage (a modest, roughly-stable
882 /// set of long-lived AI-agent clients); a public-internet-facing server
883 /// churning through unbounded distinct clients would leak memory here
884 /// with no built-in reaping mechanism.
885 sessions: Arc<Mutex<HashMap<String, StoredClientInfo>>>,
886 /// Senders for every currently-connected `GET /mcp` SSE client, pushed to
887 /// by [`Self::notify`]. `Arc<Mutex<_>>` so every clone of this `McpServer`
888 /// broadcasts to the same set of listeners.
889 ///
890 /// Entries for clients that disconnected (or whose buffer filled up) are
891 /// only pruned lazily, the next time [`Self::notify`] is called and its
892 /// `try_send` fails — not proactively, since nothing else observes the
893 /// underlying `Receiver` closing. A server that never calls `notify`
894 /// after clients disconnect will accumulate dead entries here.
895 sse_clients: Arc<Mutex<Vec<SseClient>>>,
896 /// Whether `initialize`'s advertised `capabilities` includes `"logging":{}`.
897 /// Set via [`Self::logging_enabled`]. This only controls what's
898 /// advertised — [`Self::log`] works regardless, same as [`Self::notify`]
899 /// does; a spec-honest client just wouldn't call `logging/setLevel` in
900 /// the first place if the capability was never advertised.
901 logging_enabled: bool,
902 /// The minimum [`LogLevel`] that [`Self::log`] will actually push,
903 /// settable at runtime by a client's `logging/setLevel` request. Starts
904 /// at [`LogLevel::Debug`] (the least restrictive level, i.e. nothing is
905 /// filtered) until a client requests otherwise.
906 min_log_level: Arc<Mutex<LogLevel>>,
907 /// In-flight `tools/call` cancellation flags, keyed by the request's raw
908 /// `id` JSON token. Populated just before dispatching a `tools/call` and
909 /// removed again once it returns (success, error, or the flag was never
910 /// checked) — unlike `sessions`/`sse_clients`, this map's entries are
911 /// always short-lived by construction, not merely "not evicted."
912 /// `notifications/cancelled` looks up `params.requestId` here and flips
913 /// the flag; see [`McpContext::is_cancelled`] for how (and whether) a
914 /// handler ever observes it.
915 cancellations: Arc<Mutex<HashMap<String, Arc<std::sync::atomic::AtomicBool>>>>,
916 /// Resource URI -> session ids subscribed to it via
917 /// `resources/subscribe`, consulted by [`Self::notify_resource_updated`].
918 /// A session id is only meaningful if the client also opened `GET /mcp`
919 /// with the same `Mcp-Session-Id` header (see [`SseClient`]) — a session
920 /// that only ever calls `resources/subscribe` over POST with no matching
921 /// SSE connection is subscribed in name only, since there's no channel
922 /// to push the update over. Entries are removed by explicit
923 /// `resources/unsubscribe` calls (URIs with zero subscribers are pruned
924 /// entirely); a session that disconnects without unsubscribing leaves a
925 /// harmless stale session id behind, same "no proactive eviction"
926 /// tradeoff already documented for `sessions`.
927 subscriptions: Arc<Mutex<HashMap<String, Vec<String>>>>,
928 /// In-flight `sampling/createMessage` requests this server has sent to
929 /// a client, keyed by the raw (already-quoted) request id token this
930 /// server minted. A JSON-RPC *response* has no `method` field, so
931 /// [`Self::handle_request_with_context`]/[`Self::handle_batch`]
932 /// special-case a method-less body with a recognized `id` as a
933 /// sampling reply (via [`Self::try_deliver_sampling_response`]) rather
934 /// than an invalid request. See [`McpContext::sample`] for the
935 /// send-and-block side.
936 pending_replies: Arc<Mutex<HashMap<String, mpsc::Sender<Result<String, String>>>>>,
937}
938
939/// One `GET /mcp` SSE client: its outbound channel plus the
940/// `Mcp-Session-Id` it connected with, if any. Broadcast notifications
941/// (`.notify()`, `.log()`, `list_changed`) go to every client regardless of
942/// `session_id`; [`McpServer::notify_resource_updated`] targets only
943/// clients whose `session_id` is subscribed to the changed URI.
944#[derive(Debug)]
945struct SseClient {
946 session_id: Option<String>,
947 sender: SseSender,
948}
949
950/// One `GET /mcp` SSE client's outbound channel. Bounded so a slow or stuck
951/// client can't grow memory without limit; [`McpServer::notify`] uses
952/// `try_send` (never blocks) and drops any client whose buffer is full.
953type SseSender = SyncSender<Vec<u8>>;
954
955/// Max buffered-but-unread SSE frames per client before it's treated as dead.
956const SSE_CHANNEL_CAPACITY: usize = 32;
957
958/// How often an idle SSE connection gets a `: keep-alive` comment.
959const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
960
961/// Max `completion/complete` values returned in one response, per spec
962/// guidance that servers SHOULD NOT return more than 100. A handler
963/// returning more has the rest reported via `hasMore`/`total` rather than
964/// silently included.
965const MAX_COMPLETION_VALUES: usize = 100;
966
967impl McpServer {
968 /// Create a new `McpServer`. The default MCP endpoint is `POST /mcp`.
969 pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
970 McpServer {
971 server_name: name.into(),
972 server_version: version.into(),
973 path: "/mcp".to_string(),
974 tools: Arc::new(RwLock::new(Vec::new())),
975 #[cfg(feature = "http2")]
976 async_tools: Arc::new(RwLock::new(Vec::new())),
977 resources: Arc::new(RwLock::new(Vec::new())),
978 prompts: Arc::new(RwLock::new(Vec::new())),
979 completions: Arc::new(RwLock::new(Vec::new())),
980 fallback: None,
981 auth_token: None,
982 #[cfg(feature = "sso")]
983 oauth: None,
984 page_size: None,
985 sessions: Arc::new(Mutex::new(HashMap::new())),
986 sse_clients: Arc::new(Mutex::new(Vec::new())),
987 logging_enabled: false,
988 min_log_level: Arc::new(Mutex::new(LogLevel::Debug)),
989 cancellations: Arc::new(Mutex::new(HashMap::new())),
990 subscriptions: Arc::new(Mutex::new(HashMap::new())),
991 pending_replies: Arc::new(Mutex::new(HashMap::new())),
992 }
993 }
994
995 /// Cap `tools/list`, `resources/list`, and `prompts/list` to at most `n`
996 /// items per response, enabling cursor-based pagination: a response with
997 /// more items remaining includes `"nextCursor"`, an opaque string the
998 /// client echoes back as `params.cursor` on its next call to get the next
999 /// page. `n` is clamped to a minimum of `1`.
1000 ///
1001 /// Without calling this, every registered tool/resource/prompt is
1002 /// returned in a single response — the default, and the only behavior
1003 /// before pagination existed.
1004 ///
1005 /// ```rust
1006 /// use rust_web_server::mcp::McpServer;
1007 ///
1008 /// let server = McpServer::new("my-server", "1.0").page_size(50);
1009 /// ```
1010 pub fn page_size(mut self, n: usize) -> Self {
1011 self.page_size = Some(n.max(1));
1012 self
1013 }
1014
1015 /// Push a JSON-RPC notification (no `id` — fire-and-forget, per the
1016 /// spec) to every client currently connected to the `GET /mcp` SSE
1017 /// stream, framed as an SSE `data:` event.
1018 ///
1019 /// `params_json`, if given, must already be a valid JSON value (usually
1020 /// an object) — it's spliced in verbatim, not escaped or re-serialized.
1021 ///
1022 /// Never blocks: a client whose event buffer is full (not reading fast
1023 /// enough) is treated the same as a disconnected one and dropped from
1024 /// the broadcast list, same as `notify` would drop it anyway.
1025 ///
1026 /// ```rust
1027 /// use rust_web_server::mcp::McpServer;
1028 ///
1029 /// let server = McpServer::new("my-server", "1.0");
1030 /// server.notify("notifications/message", Some(r#"{"level":"info","data":"hello"}"#));
1031 /// ```
1032 pub fn notify(&self, method: &str, params_json: Option<&str>) {
1033 let json = render_notification(method, params_json);
1034 broadcast_sse_to(&self.sse_clients, &json);
1035 }
1036
1037 /// Push `notifications/resources/updated` for `uri` to every session
1038 /// that called `resources/subscribe` for it — the mechanism behind
1039 /// live-updating resource panels (e.g. Claude Desktop watching a
1040 /// config file or a dashboard resource for changes). Call this from
1041 /// wherever your application actually changes the underlying data a
1042 /// resource represents (a file watcher, a webhook handler, a database
1043 /// trigger poll, ...).
1044 ///
1045 /// A no-op if nobody has subscribed to `uri` — including if every
1046 /// subscriber's `GET /mcp` SSE connection has since disconnected
1047 /// (pruned the same way [`Self::notify`] prunes dead broadcast clients,
1048 /// but the `subscriptions` bookkeeping for `uri` itself is untouched
1049 /// either way; only [`Self::do_resource_unsubscribe`] removes that).
1050 ///
1051 /// ```rust
1052 /// use rust_web_server::mcp::McpServer;
1053 ///
1054 /// let server = McpServer::new("my-server", "1.0");
1055 /// // Elsewhere, e.g. after reloading a watched config file:
1056 /// server.notify_resource_updated("config://main");
1057 /// ```
1058 pub fn notify_resource_updated(&self, uri: &str) {
1059 let session_ids = match self.subscriptions.lock().unwrap().get(uri) {
1060 Some(ids) if !ids.is_empty() => ids.clone(),
1061 _ => return,
1062 };
1063 let params = format!(r#"{{"uri":"{}"}}"#, json_escape(uri));
1064 let json = render_notification("notifications/resources/updated", Some(¶ms));
1065 send_sse_to_sessions(&self.sse_clients, &session_ids, &json);
1066 }
1067
1068 /// Handle `GET /mcp`: register a new SSE client and return a
1069 /// `text/event-stream` response that streams from its channel until the
1070 /// connection closes. See the module docs' SSE section for the wire
1071 /// details (keep-alive interval, backpressure behavior).
1072 ///
1073 /// Tags the client with this request's `Mcp-Session-Id` header, if
1074 /// present, so [`Self::notify_resource_updated`] can route
1075 /// `notifications/resources/updated` to just the sessions that
1076 /// subscribed via `resources/subscribe` — a client that connects
1077 /// without the header can still receive broadcasts (`.notify()`,
1078 /// `.log()`, `list_changed`) but can never be a subscription target.
1079 fn start_sse_stream(&self, request: &Request) -> Response {
1080 let session_id = request.get_header("Mcp-Session-Id".to_string()).map(|h| h.value.clone());
1081 let (tx, rx) = mpsc::sync_channel::<Vec<u8>>(SSE_CHANNEL_CAPACITY);
1082 self.sse_clients.lock().unwrap().push(SseClient { session_id, sender: tx });
1083
1084 let mut response = Response::new();
1085 response.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
1086 response.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
1087 response.headers.push(Header {
1088 name: Header::_CONTENT_TYPE.to_string(),
1089 value: "text/event-stream".to_string(),
1090 });
1091 response.headers.push(Header {
1092 name: Header::_CACHE_CONTROL.to_string(),
1093 value: "no-cache".to_string(),
1094 });
1095 response.headers.push(Header {
1096 name: "X-Accel-Buffering".to_string(),
1097 value: "no".to_string(),
1098 });
1099 response.stream_pipe = Some(Box::new(SseChannelReader::new(rx)));
1100 response
1101 }
1102
1103 /// Advertise the `logging` capability (`"logging":{}`) in `initialize`'s
1104 /// response, so clients know they can call `logging/setLevel` and expect
1105 /// `notifications/message` log entries over the `GET /mcp` SSE stream.
1106 ///
1107 /// This only changes what's *advertised* — [`Self::log`] pushes log
1108 /// entries regardless of whether this was called, same as [`Self::notify`]
1109 /// works unconditionally. A spec-honest client simply wouldn't send
1110 /// `logging/setLevel` in the first place without seeing the capability.
1111 ///
1112 /// ```rust
1113 /// use rust_web_server::mcp::McpServer;
1114 ///
1115 /// let server = McpServer::new("my-server", "1.0").logging_enabled();
1116 /// ```
1117 pub fn logging_enabled(mut self) -> Self {
1118 self.logging_enabled = true;
1119 self
1120 }
1121
1122 /// Push a `notifications/message` log entry to every client connected to
1123 /// the `GET /mcp` SSE stream, if `level` is at or above the level most
1124 /// recently set via a client's `logging/setLevel` request (or
1125 /// [`LogLevel::Debug`] — i.e. every level — if none has been set yet).
1126 ///
1127 /// `data_json` must already be a valid JSON value (the spec allows any
1128 /// type here, not just an object — a plain string is fine) — it's
1129 /// spliced in verbatim, not escaped or re-serialized. `logger`, if
1130 /// given, identifies the log's source (e.g. a module or subsystem name)
1131 /// and is escaped automatically.
1132 ///
1133 /// ```rust
1134 /// use rust_web_server::mcp::{LogLevel, McpServer};
1135 ///
1136 /// let server = McpServer::new("my-server", "1.0").logging_enabled();
1137 /// server.log(LogLevel::Warning, Some("database"), r#""connection pool exhausted""#);
1138 /// ```
1139 pub fn log(&self, level: LogLevel, logger: Option<&str>, data_json: &str) {
1140 if level < *self.min_log_level.lock().unwrap() {
1141 return;
1142 }
1143 let logger_field = match logger {
1144 Some(l) => format!(r#","logger":"{}""#, json_escape(l)),
1145 None => String::new(),
1146 };
1147 let params = format!(r#"{{"level":"{}"{logger_field},"data":{data_json}}}"#, level.as_str());
1148 self.notify("notifications/message", Some(¶ms));
1149 }
1150
1151 /// Handle `logging/setLevel`: store the requested minimum level so
1152 /// subsequent [`Self::log`] calls filter against it. Returns
1153 /// `INVALID_PARAMS` for a missing or unrecognized `params.level`.
1154 fn do_set_log_level(&self, body: &str) -> Result<String, (i32, String)> {
1155 let params = json_rpc::extract_raw(body, "params")
1156 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
1157 let level_str = json_rpc::extract_str(¶ms, "level")
1158 .ok_or((json_rpc::INVALID_PARAMS, "Missing level".to_string()))?;
1159 let level = LogLevel::parse(&level_str)
1160 .ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown log level: {level_str}")))?;
1161 *self.min_log_level.lock().unwrap() = level;
1162 Ok("{}".to_string())
1163 }
1164
1165 // ── dynamic registration ──────────────────────────────────────────────────
1166 //
1167 // Unlike `.tool()`/`.resource()`/`.prompt()` (consuming builders, called
1168 // before the server starts serving requests), these take `&self` and can
1169 // be called at any time from any thread holding a clone of this
1170 // `McpServer` — e.g. after discovering a plugin, connecting to a
1171 // database, or reacting to a hot-reloaded config file. Every clone
1172 // shares the same underlying `Arc<RwLock<Vec<_>>>`, so a mutation made
1173 // through one clone is immediately visible to every other clone,
1174 // including the ones handling concurrent requests on other threads.
1175 //
1176 // Each registration/removal pushes the corresponding
1177 // `notifications/{tools,resources,prompts}/list_changed` event (no
1178 // params, per spec) to every `GET /mcp` SSE client via `.notify()`.
1179
1180 /// Register a callable tool at runtime, exactly like [`Self::tool`] but
1181 /// without needing to own the server (and usable after it's already
1182 /// serving requests). Pushes `notifications/tools/list_changed`.
1183 ///
1184 /// ```rust
1185 /// use rust_web_server::mcp::{McpContent, McpServer};
1186 ///
1187 /// let server = McpServer::new("my-server", "1.0");
1188 ///
1189 /// // Later, from any thread holding a clone of `server`:
1190 /// server.register_tool("refresh_cache", "Reload the in-memory cache", "{}", |_args| {
1191 /// Ok(McpContent::text("cache refreshed"))
1192 /// });
1193 /// let existed = server.remove_tool("refresh_cache");
1194 /// assert!(existed);
1195 /// ```
1196 pub fn register_tool<F>(&self, name: &str, description: &str, input_schema: &str, handler: F)
1197 where
1198 F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
1199 {
1200 self.tools.write().unwrap().push(ToolDef {
1201 name: name.to_string(),
1202 description: description.to_string(),
1203 input_schema: input_schema.to_string(),
1204 annotations: None,
1205 handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
1206 });
1207 self.notify("notifications/tools/list_changed", None);
1208 }
1209
1210 /// Register an `async fn` tool handler at runtime, exactly like
1211 /// [`Self::async_tool`] but usable after the server is already serving
1212 /// requests. Requires the `http2` feature — see `async_tool`'s docs for
1213 /// the async-to-sync bridging details.
1214 #[cfg(feature = "http2")]
1215 pub fn register_async_tool<F, Fut>(&self, name: &str, description: &str, input_schema: &str, handler: F)
1216 where
1217 F: Fn(&str) -> Fut + Send + Sync + 'static,
1218 Fut: Future<Output = Result<McpContent, String>> + Send + 'static,
1219 {
1220 self.async_tools.write().unwrap().push(AsyncToolDef {
1221 name: name.to_string(),
1222 description: description.to_string(),
1223 input_schema: input_schema.to_string(),
1224 annotations: None,
1225 handler: Arc::new(move |args: String| Box::pin(handler(&args))),
1226 });
1227 self.notify("notifications/tools/list_changed", None);
1228 }
1229
1230 /// Remove a previously-registered tool by name — a sync tool registered
1231 /// via [`Self::tool`]/[`Self::register_tool`], or (with the `http2`
1232 /// feature) an async one via [`Self::async_tool`]/[`Self::register_async_tool`];
1233 /// both collections are checked. Returns `true` if a tool with that name
1234 /// existed in either and was removed. Pushes `notifications/tools/list_changed`
1235 /// only when something was actually removed.
1236 pub fn remove_tool(&self, name: &str) -> bool {
1237 #[cfg_attr(not(feature = "http2"), allow(unused_mut))]
1238 let mut removed = {
1239 let mut tools = self.tools.write().unwrap();
1240 let before = tools.len();
1241 tools.retain(|t| t.name != name);
1242 tools.len() != before
1243 };
1244
1245 #[cfg(feature = "http2")]
1246 {
1247 let mut async_tools = self.async_tools.write().unwrap();
1248 let before = async_tools.len();
1249 async_tools.retain(|t| t.name != name);
1250 removed = removed || async_tools.len() != before;
1251 }
1252
1253 if removed {
1254 self.notify("notifications/tools/list_changed", None);
1255 }
1256 removed
1257 }
1258
1259 /// Register a readable resource at runtime, exactly like [`Self::resource`].
1260 /// Pushes `notifications/resources/list_changed`.
1261 pub fn register_resource<F>(&self, uri_template: &str, name: &str, description: &str, handler: F)
1262 where
1263 F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
1264 {
1265 self.resources.write().unwrap().push(ResourceDef {
1266 uri_template: uri_template.to_string(),
1267 name: name.to_string(),
1268 description: description.to_string(),
1269 handler: Arc::new(handler),
1270 });
1271 self.notify("notifications/resources/list_changed", None);
1272 }
1273
1274 /// Remove a previously-registered resource by its exact `uri_template`
1275 /// (the same string passed to [`Self::register_resource`]/[`Self::resource`],
1276 /// not a concrete URI). Returns `true` if it existed. Pushes
1277 /// `notifications/resources/list_changed` only when something was removed.
1278 pub fn remove_resource(&self, uri_template: &str) -> bool {
1279 let removed = {
1280 let mut resources = self.resources.write().unwrap();
1281 let before = resources.len();
1282 resources.retain(|r| r.uri_template != uri_template);
1283 resources.len() != before
1284 };
1285 if removed {
1286 self.notify("notifications/resources/list_changed", None);
1287 }
1288 removed
1289 }
1290
1291 /// Register a prompt template at runtime, exactly like [`Self::prompt`]
1292 /// (no argument definitions — use [`Self::remove_prompt`] +
1293 /// [`Self::register_prompt`] if you need to change a prompt's arguments
1294 /// later; there is no dynamic equivalent of [`Self::prompt_with_args`]).
1295 /// Pushes `notifications/prompts/list_changed`.
1296 pub fn register_prompt<F>(&self, name: &str, description: &str, handler: F)
1297 where
1298 F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
1299 {
1300 self.prompts.write().unwrap().push(PromptDef {
1301 name: name.to_string(),
1302 description: description.to_string(),
1303 arguments: vec![],
1304 handler: Arc::new(handler),
1305 });
1306 self.notify("notifications/prompts/list_changed", None);
1307 }
1308
1309 /// Remove a previously-registered prompt by name. Returns `true` if it
1310 /// existed. Pushes `notifications/prompts/list_changed` only when
1311 /// something was removed.
1312 pub fn remove_prompt(&self, name: &str) -> bool {
1313 let removed = {
1314 let mut prompts = self.prompts.write().unwrap();
1315 let before = prompts.len();
1316 prompts.retain(|p| p.name != name);
1317 prompts.len() != before
1318 };
1319 if removed {
1320 self.notify("notifications/prompts/list_changed", None);
1321 }
1322 removed
1323 }
1324
1325 /// Require a bearer token on every request to the MCP endpoint.
1326 ///
1327 /// The client must send `Authorization: Bearer <token>`. Requests with a
1328 /// missing or wrong token receive `401 Unauthorized` before any JSON-RPC
1329 /// processing occurs.
1330 ///
1331 /// Store the token in an environment variable — never hard-code it:
1332 ///
1333 /// ```rust,no_run
1334 /// use rust_web_server::app::App;
1335 /// use rust_web_server::core::New;
1336 ///
1337 /// let app = App::new()
1338 /// .mcp("my-server", "1.0")
1339 /// .require_bearer(std::env::var("MCP_TOKEN").expect("MCP_TOKEN not set"));
1340 /// ```
1341 ///
1342 /// Claude Desktop config:
1343 /// ```json
1344 /// { "mcpServers": { "my-server": {
1345 /// "url": "http://localhost:7878/mcp",
1346 /// "headers": { "Authorization": "Bearer <token>" }
1347 /// }}}
1348 /// ```
1349 pub fn require_bearer(mut self, token: impl Into<String>) -> Self {
1350 self.auth_token = Some(token.into());
1351 self
1352 }
1353
1354 /// Require an OAuth 2.0 / OIDC bearer JWT — verified against a live
1355 /// JWKS endpoint — on every request to the MCP endpoint (MCP
1356 /// 2025-03-26's OAuth 2.0 authorization requirement), instead of one
1357 /// static shared secret. Requires the `sso` feature.
1358 ///
1359 /// Reuses [`crate::sso::jwks::JwksCache`] directly — the "leverage
1360 /// point" this feature was scoped around — rather than reimplementing
1361 /// JWT/JWKS verification. `provider` supplies the issuer, JWKS URI, and
1362 /// (for [`Self::execute`]'s `GET /.well-known/oauth-authorization-server`
1363 /// response) the authorization/token endpoints; use a preset
1364 /// (`OidcProvider::google()`, etc.) or `OidcProvider::discover(issuer)?`.
1365 ///
1366 /// A request with a missing, malformed, or invalid-signature bearer
1367 /// token gets `401` with `WWW-Authenticate: Bearer` before any JSON-RPC
1368 /// processing occurs — matching [`Self::require_bearer`]'s behavior. On
1369 /// success, the verified claims are available to `.tool_with_context()`
1370 /// handlers as JSON via [`McpContext::auth_claims`].
1371 ///
1372 /// If both this and [`Self::require_bearer`] are configured, OAuth
1373 /// verification takes precedence and the static token is never checked
1374 /// — configuring both is not a supported combination, just an
1375 /// unambiguous tie-break for an unusual configuration.
1376 ///
1377 /// ```rust,no_run
1378 /// use rust_web_server::app::App;
1379 /// use rust_web_server::core::New;
1380 /// use rust_web_server::sso::OidcProvider;
1381 ///
1382 /// let app = App::new()
1383 /// .mcp("my-server", "1.0")
1384 /// .require_oauth(OidcProvider::google(), "my-mcp-client-id");
1385 /// ```
1386 #[cfg(feature = "sso")]
1387 pub fn require_oauth(mut self, provider: crate::sso::OidcProvider, audience: impl Into<String>) -> Self {
1388 self.oauth = Some(Arc::new(OAuthConfig {
1389 jwks: crate::sso::jwks::JwksCache::new(&provider.jwks_uri),
1390 provider,
1391 audience: audience.into(),
1392 }));
1393 self
1394 }
1395
1396 /// Wrap an existing [`Application`] so that non-MCP requests are forwarded
1397 /// to it instead of the built-in [`App`].
1398 ///
1399 /// Use this when your existing server has custom routes, state, or
1400 /// middleware that you want to keep alongside the MCP endpoint:
1401 ///
1402 /// ```rust,no_run
1403 /// use rust_web_server::app::App;
1404 /// use rust_web_server::mcp::{McpServer, McpContent};
1405 /// use rust_web_server::response::{Response, STATUS_CODE_REASON_PHRASE};
1406 /// use rust_web_server::test_client::TestClient;
1407 ///
1408 /// let existing_app = App::with_state(42u32)
1409 /// .get("/api/hello", |_req, _params, _conn, _state| {
1410 /// Response::get_response(&STATUS_CODE_REASON_PHRASE.n200_ok, None, None)
1411 /// });
1412 ///
1413 /// let server = McpServer::new("my-app", "1.0")
1414 /// .tool("ping", "Ping", "{}", |_| Ok(McpContent::text("pong")))
1415 /// .wrap(existing_app);
1416 ///
1417 /// // Both /mcp and /api/hello are now handled by the same server.
1418 /// let client = TestClient::new(server);
1419 /// ```
1420 pub fn wrap(mut self, app: impl Application + Send + Sync + 'static) -> Self {
1421 self.fallback = Some(Arc::new(app));
1422 self
1423 }
1424
1425 /// Override the HTTP path for the MCP endpoint (default `"/mcp"`).
1426 pub fn at(mut self, path: impl Into<String>) -> Self {
1427 self.path = path.into();
1428 self
1429 }
1430
1431 /// Register a callable tool.
1432 ///
1433 /// - `name` — tool identifier (snake_case recommended)
1434 /// - `description` — human-readable description shown to the AI
1435 /// - `input_schema` — JSON Schema object for the tool's arguments
1436 /// - `handler` — closure receiving the raw `arguments` JSON string
1437 ///
1438 /// The handler returns [`McpContent`] on success or an error string. An
1439 /// error is returned to the client as `isError: true` (not a protocol error).
1440 ///
1441 /// Use [`Self::tool_with_context`] instead if the handler needs the
1442 /// caller's identity, session, or headers.
1443 pub fn tool<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
1444 where
1445 F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
1446 {
1447 self.tools.write().unwrap().push(ToolDef {
1448 name: name.to_string(),
1449 description: description.to_string(),
1450 input_schema: input_schema.to_string(),
1451 annotations: None,
1452 handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
1453 });
1454 self
1455 }
1456
1457 /// Register an `async fn` tool handler — for a tool whose work is
1458 /// naturally async (an `AsyncClient` HTTP call, an async database
1459 /// query, awaiting another future) rather than forcing it through a
1460 /// blocking call inside a plain [`Self::tool`] handler.
1461 ///
1462 /// Requires the `http2` feature (tokio). `tools/call` bridges into the
1463 /// handler's future via [`crate::async_bridge::block_on_isolated`] —
1464 /// the same mechanism [`crate::proxy::H2ReverseProxy`] and
1465 /// [`crate::async_state::AsyncAppWithState`] already use to call async
1466 /// code from a synchronous [`Application::execute`] — rather than
1467 /// `tokio::task::block_in_place`, which panics outside a `multi_thread`
1468 /// runtime; `block_on_isolated` works under any runtime flavor,
1469 /// including `current_thread`.
1470 ///
1471 /// Like [`Self::tool`] (not [`Self::tool_with_context`]), the handler
1472 /// only receives `arguments` — there is no async equivalent of
1473 /// `tool_with_context` or `tool_annotated` yet.
1474 ///
1475 /// ```rust,no_run
1476 /// # #[cfg(feature = "http2")]
1477 /// # {
1478 /// use rust_web_server::mcp::{McpContent, McpServer};
1479 ///
1480 /// let server = McpServer::new("my-server", "1.0")
1481 /// .async_tool("call_api", "Call an external API", "{}", |_args: &str| async move {
1482 /// Ok(McpContent::text("response"))
1483 /// });
1484 /// # }
1485 /// ```
1486 #[cfg(feature = "http2")]
1487 pub fn async_tool<F, Fut>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
1488 where
1489 F: Fn(&str) -> Fut + Send + Sync + 'static,
1490 Fut: Future<Output = Result<McpContent, String>> + Send + 'static,
1491 {
1492 self.async_tools.write().unwrap().push(AsyncToolDef {
1493 name: name.to_string(),
1494 description: description.to_string(),
1495 input_schema: input_schema.to_string(),
1496 annotations: None,
1497 handler: Arc::new(move |args: String| Box::pin(handler(&args))),
1498 });
1499 self
1500 }
1501
1502 /// Register a callable tool with [`ToolAnnotations`] — behavioral hints
1503 /// (read-only, destructive, idempotent, open-world) that MCP clients use
1504 /// to decide whether to warn or confirm before calling it. Otherwise
1505 /// identical to [`Self::tool`] — the handler still only receives
1506 /// `arguments`, not [`McpContext`] (there is currently no single builder
1507 /// combining annotations with per-request context; call [`Self::tool_with_context`]
1508 /// instead if you need context and don't need annotations).
1509 ///
1510 /// ```rust,no_run
1511 /// use rust_web_server::mcp::{McpContent, McpServer, ToolAnnotations};
1512 ///
1513 /// let server = McpServer::new("my-server", "1.0")
1514 /// .tool_annotated(
1515 /// "delete_file",
1516 /// "Delete a file from disk",
1517 /// r#"{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}"#,
1518 /// ToolAnnotations {
1519 /// destructive_hint: Some(true),
1520 /// read_only_hint: Some(false),
1521 /// idempotent_hint: Some(true), // deleting twice = deleting once
1522 /// ..Default::default()
1523 /// },
1524 /// |_args| Ok(McpContent::text("deleted")),
1525 /// );
1526 /// ```
1527 pub fn tool_annotated<F>(
1528 self,
1529 name: &str,
1530 description: &str,
1531 input_schema: &str,
1532 annotations: ToolAnnotations,
1533 handler: F,
1534 ) -> Self
1535 where
1536 F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
1537 {
1538 self.tools.write().unwrap().push(ToolDef {
1539 name: name.to_string(),
1540 description: description.to_string(),
1541 input_schema: input_schema.to_string(),
1542 annotations: Some(annotations),
1543 handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
1544 });
1545 self
1546 }
1547
1548 /// Register a callable tool whose handler also receives [`McpContext`] —
1549 /// caller identity/session info derived from this request's headers and
1550 /// whatever `clientInfo` this session sent at `initialize` time.
1551 ///
1552 /// Same `name`/`description`/`input_schema` semantics as [`Self::tool`];
1553 /// the only difference is the handler's first parameter.
1554 ///
1555 /// ```rust,no_run
1556 /// use rust_web_server::mcp::{McpContent, McpServer};
1557 ///
1558 /// let server = McpServer::new("my-server", "1.0")
1559 /// .tool_with_context(
1560 /// "whoami",
1561 /// "Report the caller's client info",
1562 /// "{}",
1563 /// |ctx, _args| {
1564 /// let name = ctx.client_name.as_deref().unwrap_or("unknown client");
1565 /// Ok(McpContent::text(format!("Called by {name}")))
1566 /// },
1567 /// );
1568 /// ```
1569 pub fn tool_with_context<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
1570 where
1571 F: Fn(McpContext, &str) -> Result<McpContent, String> + Send + Sync + 'static,
1572 {
1573 self.tools.write().unwrap().push(ToolDef {
1574 name: name.to_string(),
1575 description: description.to_string(),
1576 input_schema: input_schema.to_string(),
1577 annotations: None,
1578 handler: Arc::new(handler),
1579 });
1580 self
1581 }
1582
1583 /// Register a readable resource.
1584 ///
1585 /// `uri_template` uses `{param}` placeholders, e.g. `"user://{id}"`.
1586 /// The handler receives the full concrete URI string.
1587 pub fn resource<F>(self, uri_template: &str, name: &str, description: &str, handler: F) -> Self
1588 where
1589 F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
1590 {
1591 self.resources.write().unwrap().push(ResourceDef {
1592 uri_template: uri_template.to_string(),
1593 name: name.to_string(),
1594 description: description.to_string(),
1595 handler: Arc::new(handler),
1596 });
1597 self
1598 }
1599
1600 /// Register a prompt template.
1601 ///
1602 /// The handler receives the raw `arguments` JSON string and returns a
1603 /// list of [`PromptMessage`] values.
1604 pub fn prompt<F>(self, name: &str, description: &str, handler: F) -> Self
1605 where
1606 F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
1607 {
1608 self.prompts.write().unwrap().push(PromptDef {
1609 name: name.to_string(),
1610 description: description.to_string(),
1611 arguments: vec![],
1612 handler: Arc::new(handler),
1613 });
1614 self
1615 }
1616
1617 /// Register a prompt template with explicit argument definitions.
1618 pub fn prompt_with_args<F>(
1619 self,
1620 name: &str,
1621 description: &str,
1622 args: Vec<PromptArgDef>,
1623 handler: F,
1624 ) -> Self
1625 where
1626 F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
1627 {
1628 self.prompts.write().unwrap().push(PromptDef {
1629 name: name.to_string(),
1630 description: description.to_string(),
1631 arguments: args,
1632 handler: Arc::new(handler),
1633 });
1634 self
1635 }
1636
1637 /// Register an argument-completion provider for one named argument of a
1638 /// tool or prompt, so clients like Cursor and VS Code can offer
1639 /// autocomplete while the user fills in that argument.
1640 ///
1641 /// `ref_type` is `"tool"` or `"prompt"` — matched against the incoming
1642 /// `completion/complete` request's `ref.type` (`"ref/tool"`/`"ref/prompt"`
1643 /// on the wire) with the `"ref/"` prefix stripped. `ref_name` is the
1644 /// tool or prompt name this applies to. The handler receives the
1645 /// argument's name and whatever partial value the user has typed so
1646 /// far, and returns candidate completion strings (or an error, mapped
1647 /// to a JSON-RPC `INVALID_PARAMS` response).
1648 ///
1649 /// `initialize` advertises the `completions` capability automatically
1650 /// once at least one `.completion()` has been registered — there's no
1651 /// separate opt-in flag to remember.
1652 ///
1653 /// ```rust
1654 /// use rust_web_server::mcp::McpServer;
1655 ///
1656 /// let server = McpServer::new("my-server", "1.0")
1657 /// .completion("tool", "deploy", |arg_name, _partial| {
1658 /// match arg_name {
1659 /// "region" => Ok(vec!["us-east-1".to_string(), "eu-west-1".to_string()]),
1660 /// _ => Ok(vec![]),
1661 /// }
1662 /// });
1663 /// ```
1664 pub fn completion<F>(self, ref_type: &str, ref_name: &str, handler: F) -> Self
1665 where
1666 F: Fn(&str, &str) -> Result<Vec<String>, String> + Send + Sync + 'static,
1667 {
1668 self.completions.write().unwrap().push(CompletionDef {
1669 ref_type: ref_type.to_string(),
1670 ref_name: ref_name.to_string(),
1671 handler: Arc::new(handler),
1672 });
1673 self
1674 }
1675
1676 // ── request dispatch ──────────────────────────────────────────────────────
1677
1678 /// Process a raw JSON-RPC body and return an HTTP response.
1679 ///
1680 /// Equivalent to [`Self::handle_request_with_context`] with an empty
1681 /// [`McpContext`] — tool handlers registered via
1682 /// [`Self::tool_with_context`] will see every field as `None`. Prefer
1683 /// calling through [`Application::execute`] (i.e. actually serving HTTP
1684 /// requests) when you need real per-request context; this method exists
1685 /// for calling the JSON-RPC layer directly, e.g. in tests.
1686 pub fn handle_request(&self, body: &str) -> Response {
1687 self.handle_request_with_context(body, McpContext::default())
1688 }
1689
1690 /// Process a raw JSON-RPC body with an explicit [`McpContext`] and return
1691 /// an HTTP response. [`Self::execute`] calls this with a context built
1692 /// from the request's headers and this session's stored `clientInfo`;
1693 /// [`Self::handle_request`] calls this with an empty context.
1694 ///
1695 /// On a successful `initialize`, this mints a new session id (reusing
1696 /// [`crate::request_id::generate_request_id`]'s ID generator), records
1697 /// `params.clientInfo` under it, and returns the id in an
1698 /// `Mcp-Session-Id` response header — the client is expected to echo that
1699 /// header back on subsequent requests so later `tools/call`s in the same
1700 /// session can look their `clientInfo` back up.
1701 pub fn handle_request_with_context(&self, body: &str, ctx: McpContext) -> Response {
1702 let trimmed = body.trim_start();
1703 if trimmed.starts_with('[') {
1704 return self.handle_batch(trimmed, ctx);
1705 }
1706
1707 let id = json_rpc::extract_id(body);
1708
1709 let method = match json_rpc::extract_str(body, "method") {
1710 Some(m) => m,
1711 None => {
1712 // No `method` — either malformed, or a JSON-RPC *response*
1713 // to a sampling/createMessage request this server sent
1714 // earlier (responses never carry `method`).
1715 if let Some(id) = &id {
1716 if self.try_deliver_sampling_response(id, body) {
1717 return no_content();
1718 }
1719 }
1720 return rpc_error(None, json_rpc::INVALID_REQUEST, "Missing method");
1721 }
1722 };
1723
1724 if method == "notifications/cancelled" {
1725 self.handle_cancellation(body);
1726 return no_content();
1727 }
1728
1729 if method == "notifications/roots/list_changed" {
1730 self.invalidate_roots_cache(&ctx.session_id);
1731 return no_content();
1732 }
1733
1734 // Notifications have no `id` — acknowledge with 202 and no body.
1735 if method == "notifications/initialized" || (id.is_none() && method != "ping") {
1736 return no_content();
1737 }
1738
1739 let result = self.dispatch_with_cancellation(&method, body, ctx, &id);
1740 let id_str = id.as_deref().unwrap_or("null");
1741 let is_ok = result.is_ok();
1742
1743 let mut response = json_response(&Self::format_result(id_str, &result));
1744
1745 if method == "initialize" && is_ok {
1746 self.start_session(body, &mut response);
1747 }
1748
1749 response
1750 }
1751
1752 /// Process a JSON-RPC 2.0 batch request — a top-level JSON array of
1753 /// request objects sent in a single `POST /mcp` body, per the JSON-RPC
1754 /// batch spec that MCP inherits. Each element is dispatched exactly as
1755 /// [`Self::handle_request_with_context`] would dispatch it standalone;
1756 /// notifications (no `id`) contribute no entry to the response array,
1757 /// same as they'd get no response body outside a batch.
1758 ///
1759 /// An empty array (`[]`) is itself an invalid request per the JSON-RPC
1760 /// spec, so it gets one `Invalid Request` error object back rather than
1761 /// an empty array. A batch made up entirely of notifications produces no
1762 /// response body at all (`202 Accepted`), matching a single notification.
1763 ///
1764 /// If the batch contains a successful `initialize` call, the *first* one
1765 /// mints a session and attaches `Mcp-Session-Id` to the overall response,
1766 /// same as a standalone `initialize` would — sending more than one
1767 /// `initialize` in a batch is unusual and only the first is honored for
1768 /// session purposes, since one HTTP response can only carry one session id.
1769 fn handle_batch(&self, array_body: &str, ctx: McpContext) -> Response {
1770 let elements = json_rpc::split_array_elements(array_body);
1771 if elements.is_empty() {
1772 return rpc_error(None, json_rpc::INVALID_REQUEST, "Invalid Request");
1773 }
1774
1775 let mut parts: Vec<String> = Vec::new();
1776 let mut session_init_body: Option<String> = None;
1777
1778 for elem in &elements {
1779 let id = json_rpc::extract_id(elem);
1780
1781 let method = match json_rpc::extract_str(elem, "method") {
1782 Some(m) => m,
1783 None => {
1784 if let Some(id) = &id {
1785 if self.try_deliver_sampling_response(id, elem) {
1786 continue; // delivered; no response entry, same as any notification
1787 }
1788 }
1789 parts.push(Self::format_result(
1790 "null",
1791 &Err((json_rpc::INVALID_REQUEST, "Missing method".to_string())),
1792 ));
1793 continue;
1794 }
1795 };
1796
1797 if method == "notifications/cancelled" {
1798 self.handle_cancellation(elem);
1799 continue;
1800 }
1801
1802 if method == "notifications/roots/list_changed" {
1803 self.invalidate_roots_cache(&ctx.session_id);
1804 continue;
1805 }
1806
1807 if method == "notifications/initialized" || (id.is_none() && method != "ping") {
1808 continue;
1809 }
1810
1811 let result = self.dispatch_with_cancellation(&method, elem, ctx.clone(), &id);
1812 let id_str = id.as_deref().unwrap_or("null");
1813 let is_ok = result.is_ok();
1814
1815 if method == "initialize" && is_ok && session_init_body.is_none() {
1816 session_init_body = Some(elem.clone());
1817 }
1818
1819 parts.push(Self::format_result(id_str, &result));
1820 }
1821
1822 if parts.is_empty() {
1823 // Every element was a notification — no response body, same as a
1824 // single standalone notification.
1825 return no_content();
1826 }
1827
1828 let mut response = json_response(&format!("[{}]", parts.join(",")));
1829 if let Some(init_body) = session_init_body {
1830 self.start_session(&init_body, &mut response);
1831 }
1832 response
1833 }
1834
1835 /// Dispatch one already-parsed JSON-RPC `method` against `body` (the raw
1836 /// single-object message, whether it arrived standalone or as one element
1837 /// of a batch) and return the JSON-RPC `result` payload or an error.
1838 /// Shared by [`Self::handle_request_with_context`] and [`Self::handle_batch`]
1839 /// so the method table exists in exactly one place.
1840 fn dispatch(&self, method: &str, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
1841 match method {
1842 "initialize" => self.do_initialize(body),
1843 "ping" => Ok("{}".to_string()),
1844 "tools/list" => self.do_tools_list(body),
1845 "tools/call" => self.do_tools_call(body, ctx),
1846 "resources/list" => self.do_resources_list(body),
1847 "resources/read" => self.do_resources_read(body),
1848 "resources/subscribe" => self.do_resource_subscribe(body, ctx.session_id),
1849 "resources/unsubscribe" => self.do_resource_unsubscribe(body, ctx.session_id),
1850 "prompts/list" => self.do_prompts_list(body),
1851 "prompts/get" => self.do_prompts_get(body),
1852 "logging/setLevel" => self.do_set_log_level(body),
1853 "completion/complete" => self.do_completion(body),
1854 _ => Err((json_rpc::METHOD_NOT_FOUND, format!("Unknown method: {method}"))),
1855 }
1856 }
1857
1858 /// Wraps [`Self::dispatch`] for one request, additionally registering a
1859 /// cancellation flag for `tools/call` (keyed by `id`, guaranteed `Some`
1860 /// for `tools/call` — a notification-shaped `tools/call` with no `id`
1861 /// never reaches this far, see the notification checks in
1862 /// [`Self::handle_request_with_context`]/[`Self::handle_batch`]) so a
1863 /// later `notifications/cancelled` referencing this exact request can
1864 /// flip it. The entry is removed again once dispatch returns, whether
1865 /// the handler checked the flag or not — this map never accumulates
1866 /// stale entries the way `sessions`/`sse_clients` can.
1867 fn dispatch_with_cancellation(
1868 &self,
1869 method: &str,
1870 body: &str,
1871 ctx: McpContext,
1872 id: &Option<String>,
1873 ) -> Result<String, (i32, String)> {
1874 let Some(id_str) = (method == "tools/call").then_some(id.as_ref()).flatten() else {
1875 return self.dispatch(method, body, ctx);
1876 };
1877
1878 let flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
1879 self.cancellations.lock().unwrap().insert(id_str.clone(), flag.clone());
1880 let ctx = McpContext { cancellation: Some(flag), ..ctx };
1881 let result = self.dispatch(method, body, ctx);
1882 self.cancellations.lock().unwrap().remove(id_str);
1883 result
1884 }
1885
1886 /// Handle `notifications/cancelled`: look up `params.requestId` (the
1887 /// spec allows `string | number`, so this is matched as a raw JSON
1888 /// token, the same way [`McpContext::progress_token`] is) in
1889 /// `cancellations` and flip its flag if the request is still in flight.
1890 /// Silently does nothing for an unknown/already-finished request id —
1891 /// the target call may simply have already completed naturally, which
1892 /// isn't an error.
1893 fn handle_cancellation(&self, body: &str) {
1894 let Some(params) = json_rpc::extract_raw(body, "params") else { return };
1895 let Some(request_id) = json_rpc::extract_raw(¶ms, "requestId") else { return };
1896 if let Some(flag) = self.cancellations.lock().unwrap().get(&request_id) {
1897 flag.store(true, std::sync::atomic::Ordering::Relaxed);
1898 }
1899 }
1900
1901 /// Handle `notifications/roots/list_changed`: clear this session's
1902 /// cached `roots/list` result (if any), so the next
1903 /// [`McpContext::list_roots`] call in this session does a fresh round
1904 /// trip instead of serving stale data. Correlated purely by this
1905 /// request's own `Mcp-Session-Id` (`session_id`) — unlike
1906 /// `notifications/cancelled`, this notification carries no params of
1907 /// its own to key off of; it's about "your cache for *this connection*
1908 /// is stale," not any other request.
1909 fn invalidate_roots_cache(&self, session_id: &Option<String>) {
1910 let Some(sid) = session_id else { return };
1911 if let Some(info) = self.sessions.lock().unwrap().get_mut(sid) {
1912 info.roots = None;
1913 }
1914 }
1915
1916 /// If `body` is a JSON-RPC *response* (no `method`, by definition — see
1917 /// [`Self::handle_request_with_context`]) to a `sampling/createMessage`
1918 /// request this server sent and is still waiting on (`id` found in
1919 /// `pending_replies`), deliver it to the blocked [`McpContext::sample`]
1920 /// call and return `true`. Returns `false` for any `id` this server
1921 /// isn't tracking — the caller still needs to treat that as a genuine
1922 /// malformed/unrecognized message.
1923 fn try_deliver_sampling_response(&self, id: &str, body: &str) -> bool {
1924 let sender = match self.pending_replies.lock().unwrap().remove(id) {
1925 Some(s) => s,
1926 None => return false,
1927 };
1928 let payload = match json_rpc::extract_raw(body, "result") {
1929 Some(result) => Ok(result),
1930 None => {
1931 let message = json_rpc::extract_raw(body, "error")
1932 .and_then(|e| json_rpc::extract_str(&e, "message"))
1933 .unwrap_or_else(|| "sampling/createMessage request failed".to_string());
1934 Err(message)
1935 }
1936 };
1937 let _ = sender.send(payload); // the waiter may have already timed out and dropped rx
1938 true
1939 }
1940
1941 /// Render one JSON-RPC 2.0 response object — `{"jsonrpc":"2.0","result":...,"id":...}`
1942 /// or the `error` shape — from a dispatch result and its request's `id` (already
1943 /// rendered as a raw JSON token, e.g. `"1"`, `"\"abc\""`, or `"null"`).
1944 fn format_result(id_str: &str, result: &Result<String, (i32, String)>) -> String {
1945 match result {
1946 Ok(result_json) => format!(
1947 r#"{{"jsonrpc":"2.0","result":{result_json},"id":{id_str}}}"#
1948 ),
1949 Err((code, msg)) => {
1950 let escaped = json_escape(msg);
1951 format!(
1952 r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
1953 )
1954 }
1955 }
1956 }
1957
1958 /// Mint a new session id, record `body`'s `params.clientInfo` under it
1959 /// (logging the caller's identity), and attach the id to `response` as
1960 /// an `Mcp-Session-Id` header. Called once, from
1961 /// [`Self::handle_request_with_context`], right after a successful
1962 /// `initialize`.
1963 fn start_session(&self, body: &str, response: &mut Response) {
1964 let params = json_rpc::extract_raw(body, "params");
1965 let client_info = params.as_deref().and_then(|p| json_rpc::extract_raw(p, "clientInfo"));
1966 let (name, version) = match &client_info {
1967 Some(info) => (
1968 json_rpc::extract_str(info, "name"),
1969 json_rpc::extract_str(info, "version"),
1970 ),
1971 None => (None, None),
1972 };
1973 let capabilities = params.as_deref().and_then(|p| json_rpc::extract_raw(p, "capabilities"));
1974 let supports_sampling = capabilities.as_deref()
1975 .and_then(|c| json_rpc::extract_raw(c, "sampling"))
1976 .is_some();
1977 let supports_roots = capabilities.as_deref()
1978 .and_then(|c| json_rpc::extract_raw(c, "roots"))
1979 .is_some();
1980
1981 eprintln!(
1982 "[mcp] initialize from client {} v{}",
1983 name.as_deref().unwrap_or("unknown"),
1984 version.as_deref().unwrap_or("unknown"),
1985 );
1986
1987 let session_id = crate::request_id::generate_request_id();
1988 self.sessions
1989 .lock()
1990 .unwrap()
1991 .insert(session_id.clone(), StoredClientInfo { name, version, supports_sampling, supports_roots, roots: None });
1992
1993 response.headers.push(Header {
1994 name: "Mcp-Session-Id".to_string(),
1995 value: session_id,
1996 });
1997 }
1998
1999 // ── method handlers ───────────────────────────────────────────────────────
2000
2001 /// Handle `initialize`. Per spec, the server must inspect the client's
2002 /// requested `protocolVersion` and respond with the version it actually
2003 /// supports — allowing the client to abort the session if incompatible —
2004 /// rather than blindly echoing `PROTOCOL_VERSION` regardless of what was
2005 /// asked for.
2006 ///
2007 /// This server implements exactly one protocol version, so "negotiation"
2008 /// here means: if the client asked for that same version, confirm it;
2009 /// otherwise tell the client the version we actually speak (older *or*
2010 /// newer than what was requested), which is always the lower of the two
2011 /// — version strings are `YYYY-MM-DD` dates, so a plain string comparison
2012 /// already orders them correctly with no date parsing needed.
2013 ///
2014 /// `clientInfo` is *not* handled here — [`Self::handle_request_with_context`]
2015 /// extracts and stores it (under a freshly minted session id) after this
2016 /// returns, so it's only ever parsed out of `body` once per call.
2017 fn do_initialize(&self, body: &str) -> Result<String, (i32, String)> {
2018 let params = json_rpc::extract_raw(body, "params");
2019
2020 let client_version = params.as_deref().and_then(|p| json_rpc::extract_str(p, "protocolVersion"));
2021
2022 let negotiated_version: &str = match client_version.as_deref() {
2023 Some(v) if v < PROTOCOL_VERSION => v,
2024 _ => PROTOCOL_VERSION,
2025 };
2026
2027 let logging_cap = if self.logging_enabled { r#","logging":{}"# } else { "" };
2028 // completions is advertised iff at least one .completion() has been registered —
2029 // no separate opt-in flag needed, unlike logging: if nothing was registered,
2030 // completion/complete would just return empty results for everything anyway.
2031 let completions_cap = if self.completions.read().unwrap().is_empty() { "" } else { r#","completions":{}"# };
2032 // listChanged is always true: register_tool/remove_tool (and the resource/prompt
2033 // equivalents) are always available, unlike logging which is opt-in via
2034 // .logging_enabled(). resources.subscribe is also always true now that
2035 // resources/subscribe and resources/unsubscribe are implemented (MCP_TODO.md's TODO-14).
2036 let caps = format!(
2037 r#"{{"tools":{{"listChanged":true}},"resources":{{"subscribe":true,"listChanged":true}},"prompts":{{"listChanged":true}}{logging_cap}{completions_cap}}}"#
2038 );
2039 Ok(format!(
2040 r#"{{"protocolVersion":"{}","capabilities":{caps},"serverInfo":{{"name":"{}","version":"{}"}}}}"#,
2041 json_escape(negotiated_version),
2042 json_escape(&self.server_name),
2043 json_escape(&self.server_version),
2044 ))
2045 }
2046
2047 fn do_tools_list(&self, body: &str) -> Result<String, (i32, String)> {
2048 #[cfg_attr(not(feature = "http2"), allow(unused_mut))]
2049 let mut items: Vec<String> = self.tools.read().unwrap().iter()
2050 .map(|t| Self::render_tool_list_entry(&t.name, &t.description, &t.input_schema, t.annotations))
2051 .collect();
2052
2053 #[cfg(feature = "http2")]
2054 items.extend(self.async_tools.read().unwrap().iter()
2055 .map(|t| Self::render_tool_list_entry(&t.name, &t.description, &t.input_schema, t.annotations)));
2056
2057 let (page, next_cursor) = self.paginate(&items, body)?;
2058 Ok(format!(r#"{{"tools":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
2059 }
2060
2061 /// Render one `tools/list` entry — shared by the sync (`tools`) and
2062 /// (with the `http2` feature) async (`async_tools`) collections, which
2063 /// are otherwise identical from the client's point of view.
2064 fn render_tool_list_entry(name: &str, description: &str, input_schema: &str, annotations: Option<ToolAnnotations>) -> String {
2065 let annotations_field = match annotations {
2066 Some(a) => format!(r#","annotations":{}"#, a.to_json()),
2067 None => String::new(),
2068 };
2069 format!(
2070 r#"{{"name":"{}","description":"{}","inputSchema":{}{}}}"#,
2071 json_escape(name), json_escape(description), input_schema, annotations_field,
2072 )
2073 }
2074
2075 fn do_tools_call(&self, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
2076 let params = json_rpc::extract_raw(body, "params")
2077 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2078 let name = json_rpc::extract_str(¶ms, "name")
2079 .ok_or((json_rpc::INVALID_PARAMS, "Missing tool name".to_string()))?;
2080 let args = json_rpc::extract_raw(¶ms, "arguments")
2081 .unwrap_or_else(|| "{}".to_string());
2082
2083 // `_meta.progressToken` (string or number, per spec) — stored raw so
2084 // `McpContext::report_progress` can splice it back verbatim.
2085 let progress_token = json_rpc::extract_raw(¶ms, "_meta")
2086 .and_then(|meta| json_rpc::extract_raw(&meta, "progressToken"));
2087 let ctx = McpContext { progress_token, ..ctx };
2088
2089 let sync_handler = {
2090 let tools = self.tools.read().unwrap();
2091 tools.iter().find(|t| t.name == name).map(|t| t.handler.clone())
2092 };
2093 if let Some(handler) = sync_handler {
2094 return Ok(Self::format_tool_result(handler(ctx, &args)));
2095 }
2096
2097 #[cfg(feature = "http2")]
2098 {
2099 let async_handler = {
2100 let async_tools = self.async_tools.read().unwrap();
2101 async_tools.iter().find(|t| t.name == name).map(|t| t.handler.clone())
2102 };
2103 if let Some(handler) = async_handler {
2104 let args_owned = args.clone();
2105 let result = crate::async_bridge::block_on_isolated(move || handler(args_owned));
2106 return Ok(Self::format_tool_result(result));
2107 }
2108 }
2109
2110 Err((json_rpc::INVALID_PARAMS, format!("Unknown tool: {name}")))
2111 }
2112
2113 /// Render a tool handler's result as the spec's `tools/call` shape —
2114 /// `{"content":[...],"isError":false}` on success, or `isError:true`
2115 /// with the error message as a text content block. Shared by the sync
2116 /// and (with `http2`) async `tools/call` dispatch paths.
2117 fn format_tool_result(result: Result<McpContent, String>) -> String {
2118 match result {
2119 Ok(c) => format!(
2120 r#"{{"content":[{}],"isError":false}}"#,
2121 c.to_content_json(),
2122 ),
2123 Err(e) => {
2124 let escaped = json_escape(&e);
2125 format!(
2126 r#"{{"content":[{{"type":"text","text":"{escaped}"}}],"isError":true}}"#
2127 )
2128 }
2129 }
2130 }
2131
2132 fn do_resources_list(&self, body: &str) -> Result<String, (i32, String)> {
2133 let items: Vec<String> = self.resources.read().unwrap().iter().map(|r| {
2134 format!(
2135 r#"{{"uri":"{}","name":"{}","description":"{}","mimeType":"text/plain"}}"#,
2136 json_escape(&r.uri_template),
2137 json_escape(&r.name),
2138 json_escape(&r.description),
2139 )
2140 }).collect();
2141 let (page, next_cursor) = self.paginate(&items, body)?;
2142 Ok(format!(r#"{{"resources":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
2143 }
2144
2145 fn do_resources_read(&self, body: &str) -> Result<String, (i32, String)> {
2146 let params = json_rpc::extract_raw(body, "params")
2147 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2148 let uri = json_rpc::extract_str(¶ms, "uri")
2149 .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
2150
2151 let handler = {
2152 let resources = self.resources.read().unwrap();
2153 resources.iter().find(|r| uri_matches(&r.uri_template, &uri)).map(|r| r.handler.clone())
2154 }.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Resource not found: {uri}")))?;
2155
2156 match handler(&uri) {
2157 Ok(c) => {
2158 let text_esc = json_escape(&c.text);
2159 let uri_esc = json_escape(&uri);
2160 Ok(format!(
2161 r#"{{"contents":[{{"uri":"{uri_esc}","mimeType":"{}","text":"{text_esc}"}}]}}"#,
2162 c.mime(),
2163 ))
2164 }
2165 Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
2166 }
2167 }
2168
2169 /// Handle `resources/subscribe`: record that `session_id` (this
2170 /// request's `Mcp-Session-Id`) wants `notifications/resources/updated`
2171 /// pushed whenever [`Self::notify_resource_updated`] is called for
2172 /// `params.uri`. Requires a session id — without one there's no way to
2173 /// later match this subscription to a specific `GET /mcp` SSE
2174 /// connection, so a client that never sent `Mcp-Session-Id` gets
2175 /// `INVALID_PARAMS` rather than a subscription that can never fire.
2176 fn do_resource_subscribe(&self, body: &str, session_id: Option<String>) -> Result<String, (i32, String)> {
2177 let params = json_rpc::extract_raw(body, "params")
2178 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2179 let uri = json_rpc::extract_str(¶ms, "uri")
2180 .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
2181 let session_id = session_id.ok_or((
2182 json_rpc::INVALID_PARAMS,
2183 "resources/subscribe requires an Mcp-Session-Id header (call initialize first)".to_string(),
2184 ))?;
2185
2186 let mut subs = self.subscriptions.lock().unwrap();
2187 let subscribers = subs.entry(uri).or_default();
2188 if !subscribers.contains(&session_id) {
2189 subscribers.push(session_id);
2190 }
2191 Ok("{}".to_string())
2192 }
2193
2194 /// Handle `resources/unsubscribe`: the inverse of
2195 /// [`Self::do_resource_subscribe`]. Removing the last subscriber for a
2196 /// URI drops the URI's entry entirely rather than leaving an empty
2197 /// `Vec` behind.
2198 fn do_resource_unsubscribe(&self, body: &str, session_id: Option<String>) -> Result<String, (i32, String)> {
2199 let params = json_rpc::extract_raw(body, "params")
2200 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2201 let uri = json_rpc::extract_str(¶ms, "uri")
2202 .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
2203 let session_id = session_id.ok_or((
2204 json_rpc::INVALID_PARAMS,
2205 "resources/unsubscribe requires an Mcp-Session-Id header".to_string(),
2206 ))?;
2207
2208 let mut subs = self.subscriptions.lock().unwrap();
2209 if let Some(subscribers) = subs.get_mut(&uri) {
2210 subscribers.retain(|s| s != &session_id);
2211 if subscribers.is_empty() {
2212 subs.remove(&uri);
2213 }
2214 }
2215 Ok("{}".to_string())
2216 }
2217
2218 fn do_prompts_list(&self, body: &str) -> Result<String, (i32, String)> {
2219 let items: Vec<String> = self.prompts.read().unwrap().iter().map(|p| {
2220 let arg_defs: Vec<String> = p.arguments.iter().map(|a| {
2221 format!(
2222 r#"{{"name":"{}","description":"{}","required":{}}}"#,
2223 json_escape(&a.name),
2224 json_escape(&a.description),
2225 a.required,
2226 )
2227 }).collect();
2228 format!(
2229 r#"{{"name":"{}","description":"{}","arguments":[{}]}}"#,
2230 json_escape(&p.name),
2231 json_escape(&p.description),
2232 arg_defs.join(","),
2233 )
2234 }).collect();
2235 let (page, next_cursor) = self.paginate(&items, body)?;
2236 Ok(format!(r#"{{"prompts":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
2237 }
2238
2239 /// Slice `items` (already-rendered JSON object strings for one
2240 /// `*/list` response) according to [`Self::page_size`] and this
2241 /// request's `params.cursor`, returning the page and — if more items
2242 /// remain — the opaque `nextCursor` to embed in the response.
2243 ///
2244 /// Without a configured `page_size`, always returns every item and no
2245 /// cursor, i.e. pagination is fully opt-in.
2246 fn paginate<'a>(&self, items: &'a [String], body: &str) -> Result<(&'a [String], Option<String>), (i32, String)> {
2247 let page_size = match self.page_size {
2248 Some(n) => n,
2249 None => return Ok((items, None)),
2250 };
2251
2252 let cursor = json_rpc::extract_raw(body, "params")
2253 .and_then(|p| json_rpc::extract_str(&p, "cursor"));
2254
2255 let offset = match cursor {
2256 Some(c) => decode_cursor(&c)
2257 .ok_or((json_rpc::INVALID_PARAMS, "Invalid cursor".to_string()))?,
2258 None => 0,
2259 };
2260
2261 if offset >= items.len() {
2262 return Ok((&[], None));
2263 }
2264
2265 let end = (offset + page_size).min(items.len());
2266 let next_cursor = if end < items.len() { Some(encode_cursor(end)) } else { None };
2267 Ok((&items[offset..end], next_cursor))
2268 }
2269
2270 fn do_prompts_get(&self, body: &str) -> Result<String, (i32, String)> {
2271 let params = json_rpc::extract_raw(body, "params")
2272 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2273 let name = json_rpc::extract_str(¶ms, "name")
2274 .ok_or((json_rpc::INVALID_PARAMS, "Missing prompt name".to_string()))?;
2275 let args = json_rpc::extract_raw(¶ms, "arguments")
2276 .unwrap_or_else(|| "{}".to_string());
2277
2278 let (description, handler) = {
2279 let prompts = self.prompts.read().unwrap();
2280 prompts.iter().find(|p| p.name == name).map(|p| (p.description.clone(), p.handler.clone()))
2281 }.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown prompt: {name}")))?;
2282
2283 match handler(&args) {
2284 Ok(msgs) => {
2285 let msg_jsons: Vec<String> = msgs.iter().map(|m| m.to_json()).collect();
2286 Ok(format!(
2287 r#"{{"description":"{}","messages":[{}]}}"#,
2288 json_escape(&description),
2289 msg_jsons.join(","),
2290 ))
2291 }
2292 Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
2293 }
2294 }
2295
2296 /// Handle `completion/complete`: look up the registered
2297 /// [`CompletionDef`] matching `params.ref.type`/`params.ref.name`, call
2298 /// its handler with `params.argument.name`/`params.argument.value`, and
2299 /// render the spec's `{"completion":{"values":[...],"hasMore":...,
2300 /// "total":...}}` shape. No registered provider for the given ref/name
2301 /// (or an unrecognized `ref.type` not stripped of `"ref/"`) returns an
2302 /// empty `values` list rather than an error — matching the spec's own
2303 /// framing of completion as a best-effort hint, not a required capability
2304 /// per tool/prompt.
2305 fn do_completion(&self, body: &str) -> Result<String, (i32, String)> {
2306 let params = json_rpc::extract_raw(body, "params")
2307 .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
2308 let reference = json_rpc::extract_raw(¶ms, "ref")
2309 .ok_or((json_rpc::INVALID_PARAMS, "Missing ref".to_string()))?;
2310 let ref_type_raw = json_rpc::extract_str(&reference, "type")
2311 .ok_or((json_rpc::INVALID_PARAMS, "Missing ref.type".to_string()))?;
2312 let ref_type = ref_type_raw.strip_prefix("ref/").unwrap_or(&ref_type_raw);
2313 let ref_name = json_rpc::extract_str(&reference, "name")
2314 .ok_or((json_rpc::INVALID_PARAMS, "Missing ref.name".to_string()))?;
2315
2316 let argument = json_rpc::extract_raw(¶ms, "argument")
2317 .ok_or((json_rpc::INVALID_PARAMS, "Missing argument".to_string()))?;
2318 let arg_name = json_rpc::extract_str(&argument, "name")
2319 .ok_or((json_rpc::INVALID_PARAMS, "Missing argument.name".to_string()))?;
2320 let partial = json_rpc::extract_str(&argument, "value").unwrap_or_default();
2321
2322 let handler = {
2323 let completions = self.completions.read().unwrap();
2324 completions.iter()
2325 .find(|c| c.ref_type == ref_type && c.ref_name == ref_name)
2326 .map(|c| c.handler.clone())
2327 };
2328
2329 let values = match handler {
2330 Some(h) => h(&arg_name, &partial).map_err(|e| (json_rpc::INVALID_PARAMS, e))?,
2331 None => vec![],
2332 };
2333
2334 let total = values.len();
2335 let has_more = total > MAX_COMPLETION_VALUES;
2336 let page = if has_more { &values[..MAX_COMPLETION_VALUES] } else { &values[..] };
2337 let values_json: Vec<String> = page.iter().map(|v| format!(r#""{}""#, json_escape(v))).collect();
2338
2339 Ok(format!(
2340 r#"{{"completion":{{"values":[{}],"hasMore":{has_more},"total":{total}}}}}"#,
2341 values_json.join(","),
2342 ))
2343 }
2344
2345 /// Build the [`McpContext`] for an incoming request: the `Mcp-Session-Id`
2346 /// header, if present, plus whatever `clientInfo` was recorded for that
2347 /// session at `initialize` time (if this session is recognized).
2348 fn context_for(&self, request: &Request) -> McpContext {
2349 let session_id = request
2350 .get_header("Mcp-Session-Id".to_string())
2351 .map(|h| h.value.clone());
2352
2353 let (client_name, client_version, sampling_supported, roots_supported) = match &session_id {
2354 Some(sid) => match self.sessions.lock().unwrap().get(sid) {
2355 Some(info) => (info.name.clone(), info.version.clone(), info.supports_sampling, info.supports_roots),
2356 None => (None, None, false, false),
2357 },
2358 None => (None, None, false, false),
2359 };
2360
2361 McpContext {
2362 client_name,
2363 client_version,
2364 session_id,
2365 auth_claims: None,
2366 progress_token: None,
2367 sse_clients: Some(self.sse_clients.clone()),
2368 cancellation: None,
2369 sampling_supported,
2370 pending_replies: Some(self.pending_replies.clone()),
2371 roots_supported,
2372 sessions: Some(self.sessions.clone()),
2373 }
2374 }
2375}
2376
2377/// Render one JSON-RPC 2.0 notification (no `id` — fire-and-forget, per
2378/// spec) as an SSE `data:`-ready message body. Shared by [`McpServer::notify`]
2379/// and [`McpContext::report_progress`].
2380fn render_notification(method: &str, params_json: Option<&str>) -> String {
2381 let params_field = match params_json {
2382 Some(p) => format!(r#","params":{p}"#),
2383 None => String::new(),
2384 };
2385 format!(r#"{{"jsonrpc":"2.0","method":"{}"{}}}"#, json_escape(method), params_field)
2386}
2387
2388/// Send a raw pre-built JSON-RPC message to every client in `clients`,
2389/// pruning any whose channel is full or disconnected. Shared by
2390/// [`McpServer::notify`] and [`McpContext::report_progress`] — the latter
2391/// only has a clone of the broadcast list, not a whole `McpServer`.
2392fn broadcast_sse_to(clients: &Arc<Mutex<Vec<SseClient>>>, json: &str) {
2393 let frame = format!("data: {json}\n\n").into_bytes();
2394 let mut clients = clients.lock().unwrap();
2395 clients.retain(|c| c.sender.try_send(frame.clone()).is_ok());
2396}
2397
2398/// Send a raw pre-built JSON-RPC message only to clients whose `session_id`
2399/// is in `session_ids`, pruning any of *those* whose channel is full or
2400/// disconnected. A client with no `session_id`, or one not in
2401/// `session_ids`, is left untouched (kept, not sent to) — this is a
2402/// targeted send for [`McpServer::notify_resource_updated`], not a
2403/// broadcast like [`broadcast_sse_to`].
2404fn send_sse_to_sessions(clients: &Arc<Mutex<Vec<SseClient>>>, session_ids: &[String], json: &str) {
2405 let frame = format!("data: {json}\n\n").into_bytes();
2406 let mut clients = clients.lock().unwrap();
2407 clients.retain(|c| {
2408 let is_target = c.session_id.as_deref().is_some_and(|sid| session_ids.iter().any(|s| s == sid));
2409 if is_target {
2410 c.sender.try_send(frame.clone()).is_ok()
2411 } else {
2412 true
2413 }
2414 });
2415}
2416
2417// ── SSE channel reader ────────────────────────────────────────────────────────
2418
2419/// Adapts an `mpsc::Receiver<Vec<u8>>` of pre-framed SSE bytes into a
2420/// blocking [`std::io::Read`], so `Server::pipe_stream` (already written for
2421/// proxy passthrough streaming) can drive a `GET /mcp` SSE connection with no
2422/// changes to the server's write loop.
2423///
2424/// Blocks in [`Self::read`] until either a frame arrives, the sender side is
2425/// dropped (all `McpServer` clones gone — EOF, closing the connection), or
2426/// [`SSE_KEEPALIVE_INTERVAL`] elapses with nothing to send (writes a `:
2427/// keep-alive` comment instead, both to satisfy proxies that time out
2428/// silent connections and to surface a dead peer on the next write attempt).
2429struct SseChannelReader {
2430 rx: mpsc::Receiver<Vec<u8>>,
2431 leftover: Vec<u8>,
2432}
2433
2434impl SseChannelReader {
2435 fn new(rx: mpsc::Receiver<Vec<u8>>) -> Self {
2436 SseChannelReader { rx, leftover: Vec::new() }
2437 }
2438}
2439
2440impl std::io::Read for SseChannelReader {
2441 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
2442 if self.leftover.is_empty() {
2443 loop {
2444 match self.rx.recv_timeout(SSE_KEEPALIVE_INTERVAL) {
2445 Ok(frame) => { self.leftover = frame; break; }
2446 Err(mpsc::RecvTimeoutError::Timeout) => {
2447 self.leftover = b": keep-alive\n\n".to_vec();
2448 break;
2449 }
2450 Err(mpsc::RecvTimeoutError::Disconnected) => return Ok(0),
2451 }
2452 }
2453 }
2454
2455 let n = self.leftover.len().min(buf.len());
2456 buf[..n].copy_from_slice(&self.leftover[..n]);
2457 self.leftover.drain(..n);
2458 Ok(n)
2459 }
2460}
2461
2462// ── Application ───────────────────────────────────────────────────────────────
2463
2464impl Application for McpServer {
2465 fn execute(&self, request: &Request, connection: &ConnectionInfo) -> Result<Response, String> {
2466 #[cfg(feature = "sso")]
2467 if request.request_uri == "/.well-known/oauth-authorization-server" && request.method == "GET" {
2468 if let Some(oauth) = &self.oauth {
2469 return Ok(oauth_metadata_document(oauth));
2470 }
2471 }
2472
2473 if request.request_uri == self.path {
2474 #[cfg_attr(not(feature = "sso"), allow(unused_mut))]
2475 let mut auth_claims_json: Option<String> = None;
2476
2477 // OAuth (sso feature) takes precedence over a static bearer
2478 // token if both are somehow configured — see require_oauth's
2479 // doc comment.
2480 #[cfg(feature = "sso")]
2481 if let Some(oauth) = &self.oauth {
2482 let provided = request.headers.iter()
2483 .find(|h| h.name.eq_ignore_ascii_case("authorization"))
2484 .map(|h| h.value.as_str())
2485 .unwrap_or("");
2486 let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
2487 if bearer.is_empty() {
2488 return Ok(unauthorized());
2489 }
2490 let opts = crate::sso::jwks::VerifyOptions {
2491 audience: &oauth.audience,
2492 issuer: &oauth.provider.issuer,
2493 leeway_secs: 60,
2494 };
2495 match oauth.jwks.verify_jwt(bearer, &opts) {
2496 Ok(claims) => auth_claims_json = serde_json::to_string(&claims).ok(),
2497 Err(_) => return Ok(unauthorized()),
2498 }
2499 } else if let Some(expected) = &self.auth_token {
2500 // Check the static bearer token before processing any MCP request.
2501 let provided = request.headers.iter()
2502 .find(|h| h.name.eq_ignore_ascii_case("authorization"))
2503 .map(|h| h.value.as_str())
2504 .unwrap_or("");
2505 let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
2506 if bearer != expected.as_str() {
2507 return Ok(unauthorized());
2508 }
2509 }
2510 #[cfg(not(feature = "sso"))]
2511 if let Some(expected) = &self.auth_token {
2512 let provided = request.headers.iter()
2513 .find(|h| h.name.eq_ignore_ascii_case("authorization"))
2514 .map(|h| h.value.as_str())
2515 .unwrap_or("");
2516 let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
2517 if bearer != expected.as_str() {
2518 return Ok(unauthorized());
2519 }
2520 }
2521
2522 return Ok(match request.method.as_str() {
2523 "POST" => {
2524 let body = std::str::from_utf8(&request.body).unwrap_or("");
2525 let mut ctx = self.context_for(request);
2526 ctx.auth_claims = auth_claims_json;
2527 self.handle_request_with_context(body, ctx)
2528 }
2529 "GET" => self.start_sse_stream(request),
2530 "OPTIONS" => {
2531 // CORS preflight for browser-based MCP clients
2532 let mut r = Response::new();
2533 r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
2534 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
2535 r.headers.push(Header {
2536 name: "Allow".to_string(),
2537 value: "GET, POST, OPTIONS".to_string(),
2538 });
2539 r
2540 }
2541 _ => {
2542 let mut r = Response::new();
2543 r.status_code = *STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.status_code;
2544 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.reason_phrase.to_string();
2545 r.headers.push(Header {
2546 name: "Allow".to_string(),
2547 value: "GET, POST, OPTIONS".to_string(),
2548 });
2549 r.content_range_list = vec![Range::get_content_range(
2550 b"MCP endpoint only accepts GET (SSE) or POST".to_vec(),
2551 MimeType::TEXT_PLAIN.to_string(),
2552 )];
2553 r
2554 }
2555 });
2556 }
2557
2558 // Not an MCP path — fall through to the wrapped app (or built-in App).
2559 match &self.fallback {
2560 Some(app) => app.execute(request, connection),
2561 None => App::new().execute(request, connection),
2562 }
2563 }
2564}
2565
2566// ── public helper ─────────────────────────────────────────────────────────────
2567
2568/// Extract a string argument from a tool/prompt `arguments` JSON object.
2569///
2570/// ```rust
2571/// use rust_web_server::mcp::extract_arg;
2572/// assert_eq!(extract_arg(r#"{"text":"hello"}"#, "text").as_deref(), Some("hello"));
2573/// assert_eq!(extract_arg(r#"{}"#, "missing"), None);
2574/// ```
2575pub fn extract_arg(arguments: &str, name: &str) -> Option<String> {
2576 json_rpc::extract_str(arguments, name)
2577}
2578
2579// ── internal helpers ──────────────────────────────────────────────────────────
2580
2581fn json_response(body: &str) -> Response {
2582 let mut r = Response::new();
2583 r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
2584 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
2585 r.content_range_list = vec![Range::get_content_range(
2586 body.as_bytes().to_vec(),
2587 MimeType::APPLICATION_JSON.to_string(),
2588 )];
2589 r
2590}
2591
2592fn no_content() -> Response {
2593 let mut r = Response::new();
2594 r.status_code = *STATUS_CODE_REASON_PHRASE.n202_accepted.status_code;
2595 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n202_accepted.reason_phrase.to_string();
2596 r
2597}
2598
2599fn unauthorized() -> Response {
2600 let mut r = Response::new();
2601 r.status_code = *STATUS_CODE_REASON_PHRASE.n401_unauthorized.status_code;
2602 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n401_unauthorized.reason_phrase.to_string();
2603 r.headers.push(Header {
2604 name: "WWW-Authenticate".to_string(),
2605 value: "Bearer".to_string(),
2606 });
2607 r.content_range_list = vec![Range::get_content_range(
2608 b"Unauthorized".to_vec(),
2609 MimeType::TEXT_PLAIN.to_string(),
2610 )];
2611 r
2612}
2613
2614/// `GET /.well-known/oauth-authorization-server` — served whenever
2615/// [`McpServer::require_oauth`] is configured, per MCP 2025-03-26's OAuth
2616/// 2.0 authorization requirement. Not a full RFC 8414 document (this
2617/// server is a resource server, not the authorization server itself) —
2618/// just enough for a client to confirm the issuer, JWKS URI, and the
2619/// authorization/token endpoints already known from `provider`.
2620#[cfg(feature = "sso")]
2621fn oauth_metadata_document(oauth: &OAuthConfig) -> Response {
2622 let body = format!(
2623 r#"{{"issuer":"{}","authorization_endpoint":"{}","token_endpoint":"{}","jwks_uri":"{}","response_types_supported":["code"]}}"#,
2624 json_escape(&oauth.provider.issuer),
2625 json_escape(&oauth.provider.authorization_endpoint),
2626 json_escape(&oauth.provider.token_endpoint),
2627 json_escape(&oauth.provider.jwks_uri),
2628 );
2629 let mut r = Response::new();
2630 r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
2631 r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
2632 r.content_range_list = vec![Range::get_content_range(
2633 body.into_bytes(),
2634 MimeType::APPLICATION_JSON.to_string(),
2635 )];
2636 r
2637}
2638
2639fn rpc_error(id: Option<&str>, code: i32, message: &str) -> Response {
2640 let id_str = id.unwrap_or("null");
2641 let escaped = json_escape(message);
2642 json_response(&format!(
2643 r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
2644 ))
2645}
2646
2647pub(crate) fn json_escape(s: &str) -> String {
2648 let mut out = String::with_capacity(s.len() + 4);
2649 for ch in s.chars() {
2650 match ch {
2651 '"' => out.push_str("\\\""),
2652 '\\' => out.push_str("\\\\"),
2653 '\n' => out.push_str("\\n"),
2654 '\r' => out.push_str("\\r"),
2655 '\t' => out.push_str("\\t"),
2656 c if (c as u32) < 0x20 => { let _ = std::fmt::Write::write_fmt(&mut out, format_args!("\\u{:04x}", c as u32)); }
2657 c => out.push(c),
2658 }
2659 }
2660 out
2661}
2662
2663// ── pagination cursors ─────────────────────────────────────────────────────────
2664
2665/// Render `,"nextCursor":"..."` for a `*/list` response, or `""` if there's
2666/// no next page — spliced directly after the closing `]` of the items array.
2667fn next_cursor_json(next_cursor: &Option<String>) -> String {
2668 match next_cursor {
2669 Some(c) => format!(r#","nextCursor":"{}""#, json_escape(c)),
2670 None => String::new(),
2671 }
2672}
2673
2674/// Encode a `tools/list`/`resources/list`/`prompts/list` offset as the
2675/// opaque `nextCursor`/`params.cursor` string the MCP spec expects — just
2676/// base64 of the decimal offset, e.g. `50` → `"NTA="`. Callers only ever
2677/// treat this as opaque; the encoding is a private implementation detail of
2678/// this module, not a client-facing contract.
2679fn encode_cursor(offset: usize) -> String {
2680 base64_encode(offset.to_string().as_bytes())
2681}
2682
2683/// Decode a cursor produced by [`encode_cursor`]. Returns `None` for
2684/// anything that isn't valid base64 of a decimal `usize` — a malformed or
2685/// tampered cursor, not a crash.
2686fn decode_cursor(cursor: &str) -> Option<usize> {
2687 String::from_utf8(base64_decode(cursor)?).ok()?.parse().ok()
2688}
2689
2690const BASE64_TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
2691
2692fn base64_encode(data: &[u8]) -> String {
2693 let mut out = String::with_capacity((data.len() + 2) / 3 * 4);
2694 for chunk in data.chunks(3) {
2695 let b0 = chunk[0] as u32;
2696 let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
2697 let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
2698 let n = (b0 << 16) | (b1 << 8) | b2;
2699 out.push(BASE64_TABLE[((n >> 18) & 0x3F) as usize] as char);
2700 out.push(BASE64_TABLE[((n >> 12) & 0x3F) as usize] as char);
2701 out.push(if chunk.len() > 1 { BASE64_TABLE[((n >> 6) & 0x3F) as usize] as char } else { '=' });
2702 out.push(if chunk.len() > 2 { BASE64_TABLE[(n & 0x3F) as usize] as char } else { '=' });
2703 }
2704 out
2705}
2706
2707fn base64_decode(s: &str) -> Option<Vec<u8>> {
2708 fn sextet(c: u8) -> Option<u32> {
2709 match c {
2710 b'A'..=b'Z' => Some((c - b'A') as u32),
2711 b'a'..=b'z' => Some((c - b'a' + 26) as u32),
2712 b'0'..=b'9' => Some((c - b'0' + 52) as u32),
2713 b'+' => Some(62),
2714 b'/' => Some(63),
2715 _ => None,
2716 }
2717 }
2718
2719 let trimmed = s.trim_end_matches('=');
2720 let bytes = trimmed.as_bytes();
2721 let mut out = Vec::with_capacity(bytes.len() * 3 / 4 + 3);
2722 for chunk in bytes.chunks(4) {
2723 if chunk.len() == 1 {
2724 return None; // not a valid base64 length
2725 }
2726 let vals: Vec<u32> = chunk.iter().map(|&b| sextet(b)).collect::<Option<Vec<_>>>()?;
2727 let n = vals.iter().enumerate().fold(0u32, |acc, (i, &v)| acc | (v << (18 - 6 * i)));
2728 out.push(((n >> 16) & 0xFF) as u8);
2729 if vals.len() > 2 { out.push(((n >> 8) & 0xFF) as u8); }
2730 if vals.len() > 3 { out.push((n & 0xFF) as u8); }
2731 }
2732 Some(out)
2733}
2734
2735fn uri_matches(template: &str, uri: &str) -> bool {
2736 // Template `"user://{id}"` matches any URI starting with `"user://"`.
2737 match template.find('{') {
2738 Some(pos) => uri.starts_with(&template[..pos]),
2739 None => template == uri,
2740 }
2741}