rust_web_server/mcp/

mod.rs

//! Model Context Protocol (MCP) server — HTTP Streamable HTTP transport.
//!
//! [`McpServer`] implements [`Application`] so it can be passed directly to
//! [`Server::run`]. Unmatched requests fall through to the built-in [`App`]
//! controller chain (static files, health probes, etc.).
//!
//! # Quick start
//!
//! ```rust,no_run
//! use rust_web_server::server::Server;
//! use rust_web_server::mcp::{McpServer, McpContent, PromptMessage};
//! # fn main() {
//! let mcp = McpServer::new("my-server", "1.0")
//!     // A tool: callable by the AI, like a function
//!     .tool(
//!         "echo",
//!         "Echo text back",
//!         r#"{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}"#,
//!         |args| {
//!             let text = rust_web_server::mcp::extract_arg(args, "text")
//!                 .unwrap_or_else(|| "(nothing)".to_string());
//!             Ok(McpContent::text(text))
//!         },
//!     )
//!     // A resource: data the AI can read by URI
//!     .resource(
//!         "docs://{topic}",
//!         "Documentation",
//!         "Return documentation for a topic",
//!         |uri| Ok(McpContent::text(format!("Documentation for: {uri}"))),
//!     )
//!     // A prompt template: reusable message structures
//!     .prompt(
//!         "summarize",
//!         "Summarize the given text",
//!         |args| {
//!             let text = rust_web_server::mcp::extract_arg(args, "text")
//!                 .unwrap_or_else(|| "some text".to_string());
//!             Ok(vec![PromptMessage::user(format!("Please summarize: {text}"))])
//!         },
//!     );
//!
//! // let (listener, pool) = Server::setup().unwrap();
//! // Server::run(listener, pool, mcp);
//! # }
//! ```
//!
//! # MCP endpoint
//!
//! All JSON-RPC messages are sent as `POST /mcp` (override with [`.at()`](McpServer::at)).
//! The server implements the [MCP 2024-11-05 specification](https://spec.modelcontextprotocol.io).
//!
//! `GET /mcp` opens a Server-Sent Events stream for server → client push —
//! see [`McpServer::notify`] and the module docs' SSE section below.
//!
//! # SSE streaming transport
//!
//! A client that sends `GET /mcp` (instead of `POST`) gets back a
//! `text/event-stream` response that stays open indefinitely. Call
//! [`McpServer::notify`] from anywhere (a background thread, another request's
//! handler, ...) to push a JSON-RPC notification to every currently-connected
//! SSE client:
//!
//! ```rust,no_run
//! use rust_web_server::mcp::McpServer;
//!
//! let server = McpServer::new("my-server", "1.0");
//! // Elsewhere, e.g. after some background job finishes:
//! server.notify("notifications/message", Some(r#"{"level":"info","data":"job done"}"#));
//! ```
//!
//! Idle connections receive a `: keep-alive` SSE comment every 15 seconds so
//! intermediate proxies don't time them out; this doubles as the mechanism
//! that detects a client has disconnected (the next write attempt fails and
//! the connection is dropped). A client whose event buffer fills up (32
//! pending frames, unconsumed) is treated the same as a disconnected one and
//! dropped from the broadcast list — [`McpServer::notify`] never blocks the
//! calling thread waiting on a slow reader.
//!
//! This transport is only wired up for the plain HTTP/1.1 path
//! (`Server::run`/`Server::process`) — same scope as `Response::stream_pipe`
//! generally, which the HTTP/2 (`h2_handler`) and HTTP/3 (`h3_handler`)
//! handlers don't yet support for *any* response, not just this one.
//!
//! # Environment variables
//!
//! None — configure the server programmatically via the builder.

mod json_rpc;

#[cfg(test)]
mod tests;

use std::collections::HashMap;
#[cfg(feature = "http2")]
use std::future::Future;
#[cfg(feature = "http2")]
use std::pin::Pin;
use std::sync::mpsc::{self, SyncSender};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;

use crate::app::App;
use crate::application::Application;
use crate::core::New;
use crate::header::Header;
use crate::mime_type::MimeType;
use crate::range::Range;
use crate::request::Request;
use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
use crate::server::ConnectionInfo;

const PROTOCOL_VERSION: &str = "2024-11-05";

// ── public content types ──────────────────────────────────────────────────────

/// Content returned by tool and resource handlers.
///
/// Create with [`McpContent::text`] (plain text or JSON strings),
/// [`McpContent::json`] (marks MIME type as `application/json`),
/// [`McpContent::image`] (base64-encoded binary image data), or
/// [`McpContent::embedded`] (a resource embedded inline in a tool response).
#[derive(Clone, Debug)]
pub struct McpContent {
    /// `"text"`, `"image"`, or `"resource"`.
    pub kind: &'static str,
    /// The content string — text for `"text"`, base64 data for `"image"`,
    /// or the embedded resource's text for `"resource"`.
    pub text: String,
    /// Optional MIME type override (default `"text/plain"` for `"text"`;
    /// required in practice for `"image"`/`"resource"`, set by their
    /// constructors).
    pub mime_type: Option<String>,
    /// The resource URI — only set (and only serialized) for `"resource"`.
    pub uri: Option<String>,
}

impl McpContent {
    /// Plain-text content.
    pub fn text(s: impl Into<String>) -> Self {
        McpContent { kind: "text", text: s.into(), mime_type: None, uri: None }
    }

    /// JSON content — sets `mimeType` to `application/json`.
    pub fn json(s: impl Into<String>) -> Self {
        McpContent { kind: "text", text: s.into(), mime_type: Some("application/json".to_string()), uri: None }
    }

    /// Image content (screenshot, chart, generated art) — `data` is base64-encoded
    /// binary and `mime_type` is e.g. `"image/png"`.
    pub fn image(data: impl Into<String>, mime_type: impl Into<String>) -> Self {
        McpContent { kind: "image", text: data.into(), mime_type: Some(mime_type.into()), uri: None }
    }

    /// A resource embedded inline in a tool response, as opposed to one a
    /// client fetches separately via `resources/read`.
    pub fn embedded(uri: impl Into<String>, text: impl Into<String>, mime_type: impl Into<String>) -> Self {
        McpContent { kind: "resource", text: text.into(), mime_type: Some(mime_type.into()), uri: Some(uri.into()) }
    }

    fn to_content_json(&self) -> String {
        match self.kind {
            "image" => format!(
                r#"{{"type":"image","data":"{}","mimeType":"{}"}}"#,
                json_escape(&self.text),
                json_escape(self.mime_type.as_deref().unwrap_or("application/octet-stream")),
            ),
            "resource" => format!(
                r#"{{"type":"resource","resource":{{"uri":"{}","mimeType":"{}","text":"{}"}}}}"#,
                json_escape(self.uri.as_deref().unwrap_or("")),
                json_escape(self.mime_type.as_deref().unwrap_or("text/plain")),
                json_escape(&self.text),
            ),
            _ => format!(r#"{{"type":"text","text":"{}"}}"#, json_escape(&self.text)),
        }
    }

    fn mime(&self) -> &str {
        self.mime_type.as_deref().unwrap_or("text/plain")
    }
}

/// A single message in a prompt response.
#[derive(Clone, Debug)]
pub struct PromptMessage {
    /// `"user"` or `"assistant"`.
    pub role: &'static str,
    /// The message content.
    pub content: McpContent,
}

impl PromptMessage {
    /// Build a user-role message.
    pub fn user(text: impl Into<String>) -> Self {
        PromptMessage { role: "user", content: McpContent::text(text) }
    }

    /// Build an assistant-role message.
    pub fn assistant(text: impl Into<String>) -> Self {
        PromptMessage { role: "assistant", content: McpContent::text(text) }
    }

    fn to_json(&self) -> String {
        format!(
            r#"{{"role":"{}","content":{}}}"#,
            self.role,
            self.content.to_content_json(),
        )
    }
}

/// Argument definition for a prompt template.
#[derive(Clone)]
pub struct PromptArgDef {
    pub name: String,
    pub description: String,
    pub required: bool,
}

impl PromptArgDef {
    pub fn required(name: impl Into<String>, description: impl Into<String>) -> Self {
        PromptArgDef { name: name.into(), description: description.into(), required: true }
    }

    pub fn optional(name: impl Into<String>, description: impl Into<String>) -> Self {
        PromptArgDef { name: name.into(), description: description.into(), required: false }
    }
}

// ── McpContext ────────────────────────────────────────────────────────────────

/// Per-request context passed to tool handlers registered via
/// [`McpServer::tool_with_context`] — caller identity and session info that a
/// plain `Fn(&str) -> ...` tool handler has no way to see.
///
/// Constructed in [`McpServer::execute`] from the current request's headers
/// plus whatever `clientInfo` was recorded for this session at `initialize`
/// time (see [`McpServer::handle_request_with_context`]).
#[derive(Debug, Clone, Default)]
pub struct McpContext {
    /// `clientInfo.name` sent in this session's `initialize` call, if the
    /// client sent one and this request carries a recognized `Mcp-Session-Id`.
    pub client_name: Option<String>,
    /// `clientInfo.version` sent in this session's `initialize` call, under
    /// the same conditions as `client_name`.
    pub client_version: Option<String>,
    /// The `Mcp-Session-Id` header on this request, if present — the value
    /// the server minted and returned in the `initialize` response header
    /// for this session (see the module docs' Sessions section).
    pub session_id: Option<String>,
    /// Verified OAuth 2.0 / OIDC JWT claims as a JSON string (serialized
    /// [`crate::sso::OidcClaims`]), populated when [`McpServer::require_oauth`]
    /// (`sso` feature) is configured and this request's bearer token
    /// verified successfully. `None` if OAuth isn't configured, or for a
    /// hand-built context (e.g. `handle_request` in tests) that bypassed
    /// `McpServer::execute`'s auth check entirely.
    pub auth_claims: Option<String>,
    /// The raw JSON value of `params._meta.progressToken` from the
    /// triggering `tools/call` request, if the client sent one — a spec
    /// `string | number`, so this is stored pre-rendered (already correctly
    /// quoted if it's a string) rather than decoded, and spliced back
    /// verbatim by [`Self::report_progress`]. `None` for anything other than
    /// a `tools/call` whose caller asked for progress updates.
    pub progress_token: Option<String>,
    /// Shared handle back to the owning [`McpServer`]'s SSE broadcast list,
    /// used by [`Self::report_progress`]. Not `pub` — this is plumbing, not
    /// part of the context data a handler reads. `None` for a context built
    /// by hand (e.g. via [`McpServer::handle_request_with_context`] in a
    /// test) rather than through [`McpServer::execute`], in which case
    /// `report_progress` silently no-ops — there's no live server to
    /// broadcast through.
    sse_clients: Option<Arc<Mutex<Vec<SseClient>>>>,
    /// Shared flag flipped by a `notifications/cancelled` referencing this
    /// `tools/call`'s request id, checked by [`Self::is_cancelled`]. Not
    /// `pub` — see that method. `None` for anything other than a live
    /// `tools/call` dispatched through [`McpServer::execute`]/`handle_request_with_context`.
    cancellation: Option<Arc<std::sync::atomic::AtomicBool>>,
    /// Whether this session's `initialize` call declared
    /// `params.capabilities.sampling`, checked by [`Self::sample`] before
    /// sending a request the client never said it could answer.
    sampling_supported: bool,
    /// Shared handle to the owning [`McpServer`]'s in-flight
    /// server-initiated-request registry (`sampling/createMessage`,
    /// `roots/list`), used by [`Self::sample`]/[`Self::list_roots`]. Not
    /// `pub` — see `sse_clients` for the same "plumbing, not context data"
    /// reasoning; `None` under the same conditions.
    pending_replies: Option<Arc<Mutex<HashMap<String, mpsc::Sender<Result<String, String>>>>>>,
    /// Whether this session's `initialize` call declared
    /// `params.capabilities.roots`, checked by [`Self::list_roots`] before
    /// sending a request the client never said it could answer.
    roots_supported: bool,
    /// Shared handle to the owning [`McpServer`]'s session registry, used
    /// by [`Self::list_roots`] to read/write this session's cached roots
    /// list. Not `pub` — plumbing, not context data.
    sessions: Option<Arc<Mutex<HashMap<String, StoredClientInfo>>>>,
}

impl McpContext {
    /// Push a `notifications/progress` event over the SSE channel for this
    /// request's `progressToken`, if the client asked for progress updates
    /// (`params._meta.progressToken` on the triggering `tools/call`) and
    /// this context was built through a live [`McpServer`] (via `execute()`,
    /// not a bare `McpContext { .. }` — see the `sse_clients` field doc).
    ///
    /// Silently does nothing in either case — a handler doesn't need to
    /// branch on whether progress reporting is actually wired up before
    /// calling this; it's always safe to call.
    ///
    /// `total` and `message` are both optional, matching the spec's
    /// `notifications/progress` shape: `{"progressToken":...,"progress":...,
    /// "total":...,"message":"..."}` (with `total`/`message` omitted when not
    /// given here).
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::{McpContent, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_with_context("long_job", "Do something slow", "{}", |ctx, _args| {
    ///         ctx.report_progress(0.0, Some(100.0), Some("starting"));
    ///         // ... do work ...
    ///         ctx.report_progress(100.0, Some(100.0), Some("done"));
    ///         Ok(McpContent::text("done"))
    ///     });
    /// ```
    pub fn report_progress(&self, progress: f64, total: Option<f64>, message: Option<&str>) {
        let (Some(token), Some(sse_clients)) = (&self.progress_token, &self.sse_clients) else {
            return;
        };

        let total_field = match total {
            Some(t) => format!(r#","total":{t}"#),
            None => String::new(),
        };
        let message_field = match message {
            Some(m) => format!(r#","message":"{}""#, json_escape(m)),
            None => String::new(),
        };
        let params = format!(
            r#"{{"progressToken":{token},"progress":{progress}{total_field}{message_field}}}"#
        );
        let json = render_notification("notifications/progress", Some(&params));
        broadcast_sse_to(sse_clients, &json);
    }

    /// `true` if the client sent `notifications/cancelled` referencing this
    /// `tools/call`'s request id.
    ///
    /// Rust cannot forcibly interrupt a running synchronous closure — same
    /// limitation [`crate::timeout::with_timeout`] documents — so this is
    /// *cooperative* cancellation: nothing happens automatically. A handler
    /// that does its own iterative work (processing a batch of items, say)
    /// can check this between steps and return early; a handler that never
    /// checks it runs to completion regardless, exactly as before this
    /// existed.
    ///
    /// Always `false` if the client never sent a cancellation, if this
    /// wasn't dispatched as a `tools/call` (the only method cancellation
    /// applies to), or if `ctx` was built without a live server behind it.
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::{McpContent, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_with_context("process_batch", "Process many items", "{}", |ctx, _args| {
    ///         for i in 0..1_000_000 {
    ///             if ctx.is_cancelled() {
    ///                 return Err("cancelled by client".to_string());
    ///             }
    ///             // ... process item i ...
    ///         }
    ///         Ok(McpContent::text("done"))
    ///     });
    /// ```
    pub fn is_cancelled(&self) -> bool {
        self.cancellation.as_ref().is_some_and(|c| c.load(std::sync::atomic::Ordering::Relaxed))
    }

    /// Ask the connected client to run LLM inference ("sampling") and block
    /// until it replies or `timeout` elapses — the server-initiated half of
    /// `sampling/createMessage`. This reverses the usual request direction:
    /// the server sends a JSON-RPC *request* to the client over the `GET
    /// /mcp` SSE stream, and the client answers with its own `POST /mcp`
    /// carrying a JSON-RPC *response* (no `method`) that
    /// [`McpServer::handle_request_with_context`] recognizes and routes
    /// back here instead of treating it as an invalid request.
    ///
    /// Blocking (not `async fn`) is deliberate: tool handlers in this crate
    /// are plain synchronous closures — there is no async tool handler
    /// support yet (MCP_TODO.md's TODO-17) for this to `.await` inside of.
    /// The calling thread parks on the response channel for up to
    /// `timeout`; on a thread-pool server this ties up one worker for that
    /// long, same tradeoff `crate::timeout::with_timeout` already accepts
    /// for bounding a slow handler's *caller*.
    ///
    /// Fails fast, before sending anything, if: the client's `initialize`
    /// call never declared `capabilities.sampling` (it wouldn't know what
    /// to do with the request); this request has no session id (nothing to
    /// address the request to); or `ctx` has no live server behind it.
    /// Otherwise fails with a timeout error if the client never responds —
    /// including if it simply has no `GET /mcp` SSE connection open for
    /// this session, since there's no separate "not connected" signal.
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::{McpServer, PromptMessage, SamplingRequest};
    /// use std::time::Duration;
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_with_context("ask_llm", "Ask the connected client's model a question", "{}", |ctx, _args| {
    ///         let response = ctx.sample(
    ///             SamplingRequest {
    ///                 messages: vec![PromptMessage::user("What is 2+2?")],
    ///                 max_tokens: 100,
    ///                 system_prompt: None,
    ///             },
    ///             Duration::from_secs(30),
    ///         )?;
    ///         Ok(response.content)
    ///     });
    /// ```
    pub fn sample(&self, request: SamplingRequest, timeout: Duration) -> Result<SamplingResponse, String> {
        if !self.sampling_supported {
            return Err("the connected client did not declare sampling support in its initialize capabilities".to_string());
        }

        let messages_json: Vec<String> = request.messages.iter().map(|m| m.to_json()).collect();
        let system_prompt_field = match &request.system_prompt {
            Some(s) => format!(r#","systemPrompt":"{}""#, json_escape(s)),
            None => String::new(),
        };
        let params = format!(
            r#"{{"messages":[{}],"maxTokens":{}{system_prompt_field}}}"#,
            messages_json.join(","), request.max_tokens,
        );

        let result_json = self.send_and_wait("sampling/createMessage", Some(&params), timeout)?;
        parse_sampling_response(&result_json)
    }

    /// Ask the connected client which filesystem roots (workspace
    /// directories, mounted volumes, ...) it has access to — useful for a
    /// file-system-aware tool that should only operate within the client's
    /// declared workspace rather than the whole filesystem.
    ///
    /// Like [`Self::sample`], this reverses the usual request direction (a
    /// server-initiated `roots/list` request over SSE, answered by the
    /// client's own `POST /mcp`) and blocks the calling thread for up to
    /// `timeout` — see `sample`'s docs for why blocking is deliberate here.
    ///
    /// The result is cached per session: the first call after `initialize`
    /// (or after the client sends `notifications/roots/list_changed`, which
    /// invalidates the cache) does a live round trip; later calls in the
    /// same session return the cached list without sending anything. A tool
    /// handler can call this on every invocation without worrying about
    /// spamming the client with redundant `roots/list` requests.
    ///
    /// Fails fast, before sending anything, under the same conditions as
    /// [`Self::sample`] (translated to roots): the client never declared
    /// `capabilities.roots`, this request has no session id, or `ctx` has
    /// no live server behind it.
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::McpServer;
    /// use std::time::Duration;
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_with_context("list_workspace_files", "List files in the client's workspace", "{}", |ctx, _args| {
    ///         let roots = ctx.list_roots(Duration::from_secs(10))?;
    ///         let names: Vec<String> = roots.iter().map(|r| r.uri.clone()).collect();
    ///         Ok(rust_web_server::mcp::McpContent::text(names.join(", ")))
    ///     });
    /// ```
    pub fn list_roots(&self, timeout: Duration) -> Result<Vec<McpRoot>, String> {
        if !self.roots_supported {
            return Err("the connected client did not declare roots support in its initialize capabilities".to_string());
        }

        if let (Some(session_id), Some(sessions)) = (&self.session_id, &self.sessions) {
            if let Some(cached) = sessions.lock().unwrap().get(session_id).and_then(|info| info.roots.clone()) {
                return Ok(cached);
            }
        }

        let result_json = self.send_and_wait("roots/list", None, timeout)?;
        let roots = parse_roots_response(&result_json)?;

        if let (Some(session_id), Some(sessions)) = (&self.session_id, &self.sessions) {
            if let Some(info) = sessions.lock().unwrap().get_mut(session_id) {
                info.roots = Some(roots.clone());
            }
        }

        Ok(roots)
    }

    /// Send a JSON-RPC request to the client for this context's session and
    /// block for the reply — the shared send-and-correlate mechanic behind
    /// [`Self::sample`] and [`Self::list_roots`]. Returns the raw `result`
    /// JSON string (not yet parsed into either method's own response type)
    /// on success.
    fn send_and_wait(&self, method: &str, params_json: Option<&str>, timeout: Duration) -> Result<String, String> {
        let session_id = self.session_id.clone()
            .ok_or_else(|| format!("{method} requires a session (Mcp-Session-Id)"))?;
        let sse_clients = self.sse_clients.as_ref()
            .ok_or_else(|| format!("{method} requires a live server (this context has none)"))?;
        let pending = self.pending_replies.as_ref()
            .ok_or_else(|| format!("{method} requires a live server (this context has none)"))?;

        let request_id = format!("\"{}\"", crate::request_id::generate_request_id());
        let (tx, rx) = mpsc::channel::<Result<String, String>>();
        pending.lock().unwrap().insert(request_id.clone(), tx);

        let params_field = match params_json {
            Some(p) => format!(r#","params":{p}"#),
            None => String::new(),
        };
        let rpc = format!(r#"{{"jsonrpc":"2.0","id":{request_id},"method":"{method}"{params_field}}}"#);

        send_sse_to_sessions(sse_clients, &[session_id], &rpc);

        let outcome = rx.recv_timeout(timeout);
        pending.lock().unwrap().remove(&request_id);

        match outcome {
            Ok(Ok(result_json)) => Ok(result_json),
            Ok(Err(e)) => Err(e),
            Err(_) => Err(format!("{method} timed out waiting for the client's response")),
        }
    }
}

/// A request to ask the connected MCP client to run LLM inference
/// ("sampling") — build one and pass it to [`McpContext::sample`].
///
/// `messages` reuses [`PromptMessage`] rather than introducing a near-identical
/// type — the spec's sampling message shape (`{"role":...,"content":{"type":"text",...}}`)
/// is exactly what `PromptMessage` already models.
pub struct SamplingRequest {
    pub messages: Vec<PromptMessage>,
    pub max_tokens: u32,
    /// An optional system prompt to steer the client's model, per the
    /// spec's `params.systemPrompt`.
    pub system_prompt: Option<String>,
}

/// The client's answer to a [`SamplingRequest`], returned by
/// [`McpContext::sample`].
pub struct SamplingResponse {
    /// Almost always `"assistant"`, per spec.
    pub role: String,
    pub content: McpContent,
    /// The model the client actually used to generate this response, e.g.
    /// `"claude-3-sonnet-20240307"`.
    pub model: String,
    /// Why the model stopped, e.g. `"endTurn"`, `"maxTokens"`, `"stopSequence"`.
    /// `None` if the client didn't report one.
    pub stop_reason: Option<String>,
}

/// Parse a `sampling/createMessage` response's `result` object (already
/// extracted from the enclosing JSON-RPC envelope) into a [`SamplingResponse`].
fn parse_sampling_response(result_json: &str) -> Result<SamplingResponse, String> {
    let role = json_rpc::extract_str(result_json, "role").unwrap_or_else(|| "assistant".to_string());
    let content_json = json_rpc::extract_raw(result_json, "content")
        .ok_or_else(|| "sampling response missing content".to_string())?;
    let kind = json_rpc::extract_str(&content_json, "type").unwrap_or_else(|| "text".to_string());
    let content = match kind.as_str() {
        "image" => McpContent::image(
            json_rpc::extract_str(&content_json, "data").unwrap_or_default(),
            json_rpc::extract_str(&content_json, "mimeType").unwrap_or_default(),
        ),
        _ => McpContent::text(json_rpc::extract_str(&content_json, "text").unwrap_or_default()),
    };
    let model = json_rpc::extract_str(result_json, "model").unwrap_or_default();
    let stop_reason = json_rpc::extract_str(result_json, "stopReason");
    Ok(SamplingResponse { role, content, model, stop_reason })
}

/// One filesystem root a client has access to, returned by
/// [`McpContext::list_roots`] — a workspace directory, a mounted volume,
/// or similar. `uri` is typically a `file://` URI per spec.
#[derive(Debug, Clone)]
pub struct McpRoot {
    pub uri: String,
    /// A human-readable label for the root, if the client provided one.
    pub name: Option<String>,
}

/// Parse a `roots/list` response's `result` object (already extracted from
/// the enclosing JSON-RPC envelope) into a `Vec<McpRoot>`.
fn parse_roots_response(result_json: &str) -> Result<Vec<McpRoot>, String> {
    let roots_array = json_rpc::extract_raw(result_json, "roots")
        .ok_or_else(|| "roots/list response missing roots".to_string())?;
    let roots = json_rpc::split_array_elements(&roots_array).iter().map(|elem| {
        McpRoot {
            uri: json_rpc::extract_str(elem, "uri").unwrap_or_default(),
            name: json_rpc::extract_str(elem, "name"),
        }
    }).collect();
    Ok(roots)
}

/// Configuration installed by [`McpServer::require_oauth`] (`sso` feature).
/// Reuses [`crate::sso::jwks::JwksCache`] and [`crate::sso::OidcProvider`]
/// directly rather than re-implementing JWKS fetch/verify — see
/// `require_oauth`'s doc comment for the full rationale.
#[cfg(feature = "sso")]
struct OAuthConfig {
    jwks: crate::sso::jwks::JwksCache,
    provider: crate::sso::OidcProvider,
    audience: String,
}

/// `clientInfo` recorded for one session at `initialize` time, looked up by
/// `Mcp-Session-Id` for later requests in the same session. See
/// `McpServer`'s `sessions` field doc comment for the unbounded-growth caveat.
#[derive(Debug, Clone, Default)]
struct StoredClientInfo {
    name: Option<String>,
    version: Option<String>,
    /// Whether this session's `initialize` call declared `params.capabilities.sampling`
    /// (spec: the client, not the server, declares sampling support — the
    /// presence of the key is what matters, not its contents, which are
    /// normally an empty object). Checked by [`McpContext::sample`] before
    /// sending a `sampling/createMessage` request that the client would
    /// have no idea how to answer.
    supports_sampling: bool,
    /// Whether this session's `initialize` call declared `params.capabilities.roots`,
    /// checked by [`McpContext::list_roots`] the same way `supports_sampling` gates `sample`.
    supports_roots: bool,
    /// Cached result of this session's last `roots/list` round trip, read
    /// and written by [`McpContext::list_roots`]. `None` means "never
    /// fetched, or invalidated" — set back to `None` when this session
    /// sends `notifications/roots/list_changed`, so the next `list_roots`
    /// call does a fresh round trip instead of serving stale data.
    roots: Option<Vec<McpRoot>>,
}

// ── ToolAnnotations ───────────────────────────────────────────────────────────

/// Behavioral hints for a tool, per the MCP 2025-03-26 spec's tool
/// annotations. Clients (Claude Desktop and others) use these to decide
/// whether to warn or ask for confirmation before calling a tool — e.g. skip
/// confirmation for a read-only tool, or warn before a destructive one.
///
/// Every field is a *hint*, not something this server enforces or verifies —
/// nothing stops a handler registered with `read_only_hint: Some(true)` from
/// actually writing to disk. A well-behaved server sets these accurately;
/// a client is free to ignore them or ask for confirmation anyway.
///
/// Register with [`McpServer::tool_annotated`]. Build one with plain struct
/// syntax — every field defaults to `None` (no hint given, the client's own
/// default applies):
///
/// ```rust
/// use rust_web_server::mcp::ToolAnnotations;
///
/// let destructive = ToolAnnotations {
///     destructive_hint: Some(true),
///     read_only_hint: Some(false),
///     ..Default::default()
/// };
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct ToolAnnotations {
    /// The tool does not modify its environment.
    pub read_only_hint: Option<bool>,
    /// The tool may perform destructive updates (only meaningful when
    /// `read_only_hint` is not `Some(true)`).
    pub destructive_hint: Option<bool>,
    /// Calling the tool repeatedly with the same arguments has no additional
    /// effect beyond the first call.
    pub idempotent_hint: Option<bool>,
    /// The tool may interact with an open-ended set of external entities
    /// (e.g. web search), as opposed to a fixed, closed set.
    pub open_world_hint: Option<bool>,
}

impl ToolAnnotations {
    /// Render as a JSON object containing only the hints that are `Some`,
    /// using the spec's camelCase key names. Returns `"{}"` if every field
    /// is `None`.
    fn to_json(self) -> String {
        let mut fields = Vec::with_capacity(4);
        if let Some(v) = self.read_only_hint {
            fields.push(format!(r#""readOnlyHint":{v}"#));
        }
        if let Some(v) = self.destructive_hint {
            fields.push(format!(r#""destructiveHint":{v}"#));
        }
        if let Some(v) = self.idempotent_hint {
            fields.push(format!(r#""idempotentHint":{v}"#));
        }
        if let Some(v) = self.open_world_hint {
            fields.push(format!(r#""openWorldHint":{v}"#));
        }
        format!("{{{}}}", fields.join(","))
    }
}

// ── LogLevel ──────────────────────────────────────────────────────────────────

/// RFC 5424 syslog severity levels, as used by the MCP `logging/setLevel`
/// request and `notifications/message` log entries — ordered from most to
/// least verbose so `level < min_level` comparisons work directly (this
/// relies on declaration order matching severity order; don't reorder the
/// variants).
///
/// ```rust
/// use rust_web_server::mcp::LogLevel;
///
/// assert!(LogLevel::Debug < LogLevel::Warning);
/// assert!(LogLevel::Emergency > LogLevel::Error);
/// assert_eq!(LogLevel::parse("warning"), Some(LogLevel::Warning));
/// assert_eq!(LogLevel::Warning.as_str(), "warning");
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum LogLevel {
    Debug,
    Info,
    Notice,
    Warning,
    Error,
    Critical,
    Alert,
    Emergency,
}

impl LogLevel {
    /// Parse the MCP spec's lowercase level name. Returns `None` for
    /// anything that isn't one of the eight recognized levels.
    pub fn parse(s: &str) -> Option<Self> {
        match s {
            "debug"     => Some(LogLevel::Debug),
            "info"      => Some(LogLevel::Info),
            "notice"    => Some(LogLevel::Notice),
            "warning"   => Some(LogLevel::Warning),
            "error"     => Some(LogLevel::Error),
            "critical"  => Some(LogLevel::Critical),
            "alert"     => Some(LogLevel::Alert),
            "emergency" => Some(LogLevel::Emergency),
            _           => None,
        }
    }

    /// The MCP spec's lowercase level name, e.g. `"warning"`.
    pub fn as_str(self) -> &'static str {
        match self {
            LogLevel::Debug     => "debug",
            LogLevel::Info      => "info",
            LogLevel::Notice    => "notice",
            LogLevel::Warning   => "warning",
            LogLevel::Error     => "error",
            LogLevel::Critical  => "critical",
            LogLevel::Alert     => "alert",
            LogLevel::Emergency => "emergency",
        }
    }
}

// ── internal handler registrations ───────────────────────────────────────────

type ToolFn     = Arc<dyn Fn(McpContext, &str) -> Result<McpContent, String>    + Send + Sync>;
type ResourceFn = Arc<dyn Fn(&str) -> Result<McpContent, String>    + Send + Sync>;
type PromptFn   = Arc<dyn Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync>;
/// `Fn(argument_name, partial_value) -> candidate completion strings`.
type CompletionFn = Arc<dyn Fn(&str, &str) -> Result<Vec<String>, String> + Send + Sync>;
/// An async tool handler, boxed so `AsyncToolDef` can store handlers with
/// different concrete `Future` types uniformly. The `Future` itself doesn't
/// need `Send` for [`crate::async_bridge::block_on_isolated`] (which drives
/// it to completion without ever moving it across a thread boundary), but
/// it's required here anyway for consistency with [`crate::async_state`]'s
/// equivalent handler type — the common case in practice, and one less
/// surprise for anyone switching between the two async handler styles.
#[cfg(feature = "http2")]
type AsyncToolFn = Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<McpContent, String>> + Send>> + Send + Sync>;

#[derive(Clone)]
struct ToolDef {
    name: String,
    description: String,
    input_schema: String,
    annotations: Option<ToolAnnotations>,
    handler: ToolFn,
}

/// A tool registered via [`McpServer::async_tool`]/[`McpServer::register_async_tool`]
/// — stored separately from [`ToolDef`] rather than unifying the two behind
/// a handler enum, since sync and async tools are otherwise identical and
/// keeping them apart touches far less of the existing (already
/// substantial) tool-handling code. `tools/list` and `tools/call` merge
/// both collections transparently; from the client's perspective there is
/// no distinction.
#[cfg(feature = "http2")]
#[derive(Clone)]
struct AsyncToolDef {
    name: String,
    description: String,
    input_schema: String,
    annotations: Option<ToolAnnotations>,
    handler: AsyncToolFn,
}

#[derive(Clone)]
struct ResourceDef {
    uri_template: String,
    name: String,
    description: String,
    handler: ResourceFn,
}

#[derive(Clone)]
struct PromptDef {
    name: String,
    description: String,
    arguments: Vec<PromptArgDef>,
    handler: PromptFn,
}

/// One `.completion()` registration — completion candidates for a single
/// named argument of a single tool or prompt. `ref_type` is the short form
/// passed to `.completion()` (e.g. `"tool"`, `"prompt"`), matched against the
/// request's `ref.type` (e.g. `"ref/tool"`) with the `"ref/"` prefix
/// stripped, not the raw wire value.
#[derive(Clone)]
struct CompletionDef {
    ref_type: String,
    ref_name: String,
    handler: CompletionFn,
}

// ── McpServer ─────────────────────────────────────────────────────────────────

/// An HTTP server that implements the MCP 2024-11-05 protocol.
///
/// Register tools, resources, and prompts with the builder methods, then pass
/// the server to [`Server::run`] (or [`Server::run_tls`]) as an [`Application`].
/// Requests that do not match the MCP endpoint fall through to the built-in
/// [`App`] controller chain.
#[derive(Clone)]
pub struct McpServer {
    server_name: String,
    server_version: String,
    path: String,
    /// `Arc<RwLock<_>>` (not a plain `Vec`) so a running server's tool list
    /// can be mutated at runtime — see [`Self::register_tool`]/[`Self::remove_tool`]
    /// — and every clone of this `McpServer` (each connection thread gets
    /// one) sees the same live list.
    tools: Arc<RwLock<Vec<ToolDef>>>,
    /// Tools registered via [`Self::async_tool`]/[`Self::register_async_tool`],
    /// bridged into synchronous [`Application::execute`] via
    /// [`crate::async_bridge::block_on_isolated`] at call time. Kept
    /// separate from `tools` rather than merged into one collection — see
    /// [`AsyncToolDef`]'s doc comment.
    #[cfg(feature = "http2")]
    async_tools: Arc<RwLock<Vec<AsyncToolDef>>>,
    resources: Arc<RwLock<Vec<ResourceDef>>>,
    prompts: Arc<RwLock<Vec<PromptDef>>>,
    /// Argument completion providers registered via [`Self::completion`].
    /// `initialize` advertises the `completions` capability iff this is
    /// non-empty at that moment.
    completions: Arc<RwLock<Vec<CompletionDef>>>,
    fallback: Option<Arc<dyn Application + Send + Sync>>,
    auth_token: Option<String>,
    /// Origins a browser may call the endpoint from (see
    /// [`Self::allowed_origins`]); a request with any other `Origin` is
    /// refused.
    allowed_origins: Vec<String>,
    /// Installed by [`Self::require_oauth`] (`sso` feature) — an alternative
    /// to `auth_token` that verifies a client-supplied OAuth 2.0 / OIDC
    /// bearer JWT via JWKS instead of comparing against one static shared
    /// secret. If both are somehow configured, OAuth verification takes
    /// precedence and `auth_token` is ignored (see [`Self::require_oauth`]).
    #[cfg(feature = "sso")]
    oauth: Option<Arc<OAuthConfig>>,
    /// Max items per page for `tools/list`/`resources/list`/`prompts/list`,
    /// set via [`Self::page_size`]. `None` (the default) means no pagination
    /// — every item comes back in one response, same as before pagination
    /// existed.
    page_size: Option<usize>,
    /// `clientInfo` recorded per `Mcp-Session-Id`, minted at `initialize` time.
    /// `Arc<Mutex<_>>` so every clone of this `McpServer` shares one map.
    ///
    /// This map only grows — nothing ever removes an entry, since there's no
    /// session-termination signal in the MCP Streamable HTTP transport to key
    /// eviction off of. Fine for the expected usage (a modest, roughly-stable
    /// set of long-lived AI-agent clients); a public-internet-facing server
    /// churning through unbounded distinct clients would leak memory here
    /// with no built-in reaping mechanism.
    sessions: Arc<Mutex<HashMap<String, StoredClientInfo>>>,
    /// Senders for every currently-connected `GET /mcp` SSE client, pushed to
    /// by [`Self::notify`]. `Arc<Mutex<_>>` so every clone of this `McpServer`
    /// broadcasts to the same set of listeners.
    ///
    /// Entries for clients that disconnected (or whose buffer filled up) are
    /// only pruned lazily, the next time [`Self::notify`] is called and its
    /// `try_send` fails — not proactively, since nothing else observes the
    /// underlying `Receiver` closing. A server that never calls `notify`
    /// after clients disconnect will accumulate dead entries here.
    sse_clients: Arc<Mutex<Vec<SseClient>>>,
    /// Whether `initialize`'s advertised `capabilities` includes `"logging":{}`.
    /// Set via [`Self::logging_enabled`]. This only controls what's
    /// advertised — [`Self::log`] works regardless, same as [`Self::notify`]
    /// does; a spec-honest client just wouldn't call `logging/setLevel` in
    /// the first place if the capability was never advertised.
    logging_enabled: bool,
    /// The minimum [`LogLevel`] that [`Self::log`] will actually push,
    /// settable at runtime by a client's `logging/setLevel` request. Starts
    /// at [`LogLevel::Debug`] (the least restrictive level, i.e. nothing is
    /// filtered) until a client requests otherwise.
    min_log_level: Arc<Mutex<LogLevel>>,
    /// In-flight `tools/call` cancellation flags, keyed by the request's raw
    /// `id` JSON token. Populated just before dispatching a `tools/call` and
    /// removed again once it returns (success, error, or the flag was never
    /// checked) — unlike `sessions`/`sse_clients`, this map's entries are
    /// always short-lived by construction, not merely "not evicted."
    /// `notifications/cancelled` looks up `params.requestId` here and flips
    /// the flag; see [`McpContext::is_cancelled`] for how (and whether) a
    /// handler ever observes it.
    cancellations: Arc<Mutex<HashMap<String, Arc<std::sync::atomic::AtomicBool>>>>,
    /// Resource URI -> session ids subscribed to it via
    /// `resources/subscribe`, consulted by [`Self::notify_resource_updated`].
    /// A session id is only meaningful if the client also opened `GET /mcp`
    /// with the same `Mcp-Session-Id` header (see [`SseClient`]) — a session
    /// that only ever calls `resources/subscribe` over POST with no matching
    /// SSE connection is subscribed in name only, since there's no channel
    /// to push the update over. Entries are removed by explicit
    /// `resources/unsubscribe` calls (URIs with zero subscribers are pruned
    /// entirely); a session that disconnects without unsubscribing leaves a
    /// harmless stale session id behind, same "no proactive eviction"
    /// tradeoff already documented for `sessions`.
    subscriptions: Arc<Mutex<HashMap<String, Vec<String>>>>,
    /// In-flight `sampling/createMessage` requests this server has sent to
    /// a client, keyed by the raw (already-quoted) request id token this
    /// server minted. A JSON-RPC *response* has no `method` field, so
    /// [`Self::handle_request_with_context`]/[`Self::handle_batch`]
    /// special-case a method-less body with a recognized `id` as a
    /// sampling reply (via [`Self::try_deliver_sampling_response`]) rather
    /// than an invalid request. See [`McpContext::sample`] for the
    /// send-and-block side.
    pending_replies: Arc<Mutex<HashMap<String, mpsc::Sender<Result<String, String>>>>>,
}

/// One `GET /mcp` SSE client: its outbound channel plus the
/// `Mcp-Session-Id` it connected with, if any. Broadcast notifications
/// (`.notify()`, `.log()`, `list_changed`) go to every client regardless of
/// `session_id`; [`McpServer::notify_resource_updated`] targets only
/// clients whose `session_id` is subscribed to the changed URI.
#[derive(Debug)]
struct SseClient {
    session_id: Option<String>,
    sender: SseSender,
}

/// One `GET /mcp` SSE client's outbound channel. Bounded so a slow or stuck
/// client can't grow memory without limit; [`McpServer::notify`] uses
/// `try_send` (never blocks) and drops any client whose buffer is full.
type SseSender = SyncSender<Vec<u8>>;

/// Max buffered-but-unread SSE frames per client before it's treated as dead.
const SSE_CHANNEL_CAPACITY: usize = 32;

/// How often an idle SSE connection gets a `: keep-alive` comment.
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);

/// Max `completion/complete` values returned in one response, per spec
/// guidance that servers SHOULD NOT return more than 100. A handler
/// returning more has the rest reported via `hasMore`/`total` rather than
/// silently included.
const MAX_COMPLETION_VALUES: usize = 100;

impl McpServer {
    /// Create a new `McpServer`.  The default MCP endpoint is `POST /mcp`.
    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
        McpServer {
            server_name: name.into(),
            server_version: version.into(),
            path: "/mcp".to_string(),
            tools: Arc::new(RwLock::new(Vec::new())),
            #[cfg(feature = "http2")]
            async_tools: Arc::new(RwLock::new(Vec::new())),
            resources: Arc::new(RwLock::new(Vec::new())),
            prompts: Arc::new(RwLock::new(Vec::new())),
            completions: Arc::new(RwLock::new(Vec::new())),
            fallback: None,
            auth_token: None,
            allowed_origins: Vec::new(),
            #[cfg(feature = "sso")]
            oauth: None,
            page_size: None,
            sessions: Arc::new(Mutex::new(HashMap::new())),
            sse_clients: Arc::new(Mutex::new(Vec::new())),
            logging_enabled: false,
            min_log_level: Arc::new(Mutex::new(LogLevel::Debug)),
            cancellations: Arc::new(Mutex::new(HashMap::new())),
            subscriptions: Arc::new(Mutex::new(HashMap::new())),
            pending_replies: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Cap `tools/list`, `resources/list`, and `prompts/list` to at most `n`
    /// items per response, enabling cursor-based pagination: a response with
    /// more items remaining includes `"nextCursor"`, an opaque string the
    /// client echoes back as `params.cursor` on its next call to get the next
    /// page. `n` is clamped to a minimum of `1`.
    ///
    /// Without calling this, every registered tool/resource/prompt is
    /// returned in a single response — the default, and the only behavior
    /// before pagination existed.
    ///
    /// ```rust
    /// use rust_web_server::mcp::McpServer;
    ///
    /// let server = McpServer::new("my-server", "1.0").page_size(50);
    /// ```
    pub fn page_size(mut self, n: usize) -> Self {
        self.page_size = Some(n.max(1));
        self
    }

    /// Push a JSON-RPC notification (no `id` — fire-and-forget, per the
    /// spec) to every client currently connected to the `GET /mcp` SSE
    /// stream, framed as an SSE `data:` event.
    ///
    /// `params_json`, if given, must already be a valid JSON value (usually
    /// an object) — it's spliced in verbatim, not escaped or re-serialized.
    ///
    /// Never blocks: a client whose event buffer is full (not reading fast
    /// enough) is treated the same as a disconnected one and dropped from
    /// the broadcast list, same as `notify` would drop it anyway.
    ///
    /// ```rust
    /// use rust_web_server::mcp::McpServer;
    ///
    /// let server = McpServer::new("my-server", "1.0");
    /// server.notify("notifications/message", Some(r#"{"level":"info","data":"hello"}"#));
    /// ```
    pub fn notify(&self, method: &str, params_json: Option<&str>) {
        let json = render_notification(method, params_json);
        broadcast_sse_to(&self.sse_clients, &json);
    }

    /// Push `notifications/resources/updated` for `uri` to every session
    /// that called `resources/subscribe` for it — the mechanism behind
    /// live-updating resource panels (e.g. Claude Desktop watching a
    /// config file or a dashboard resource for changes). Call this from
    /// wherever your application actually changes the underlying data a
    /// resource represents (a file watcher, a webhook handler, a database
    /// trigger poll, ...).
    ///
    /// A no-op if nobody has subscribed to `uri` — including if every
    /// subscriber's `GET /mcp` SSE connection has since disconnected
    /// (pruned the same way [`Self::notify`] prunes dead broadcast clients,
    /// but the `subscriptions` bookkeeping for `uri` itself is untouched
    /// either way; only [`Self::do_resource_unsubscribe`] removes that).
    ///
    /// ```rust
    /// use rust_web_server::mcp::McpServer;
    ///
    /// let server = McpServer::new("my-server", "1.0");
    /// // Elsewhere, e.g. after reloading a watched config file:
    /// server.notify_resource_updated("config://main");
    /// ```
    pub fn notify_resource_updated(&self, uri: &str) {
        let session_ids = match self.subscriptions.lock().unwrap().get(uri) {
            Some(ids) if !ids.is_empty() => ids.clone(),
            _ => return,
        };
        let params = format!(r#"{{"uri":"{}"}}"#, json_escape(uri));
        let json = render_notification("notifications/resources/updated", Some(&params));
        send_sse_to_sessions(&self.sse_clients, &session_ids, &json);
    }

    /// Handle `GET /mcp`: register a new SSE client and return a
    /// `text/event-stream` response that streams from its channel until the
    /// connection closes. See the module docs' SSE section for the wire
    /// details (keep-alive interval, backpressure behavior).
    ///
    /// Tags the client with this request's `Mcp-Session-Id` header, if
    /// present, so [`Self::notify_resource_updated`] can route
    /// `notifications/resources/updated` to just the sessions that
    /// subscribed via `resources/subscribe` — a client that connects
    /// without the header can still receive broadcasts (`.notify()`,
    /// `.log()`, `list_changed`) but can never be a subscription target.
    fn start_sse_stream(&self, request: &Request) -> Response {
        let session_id = request.get_header("Mcp-Session-Id".to_string()).map(|h| h.value.clone());
        let (tx, rx) = mpsc::sync_channel::<Vec<u8>>(SSE_CHANNEL_CAPACITY);
        self.sse_clients.lock().unwrap().push(SseClient { session_id, sender: tx });

        let mut response = Response::new();
        response.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
        response.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
        response.headers.push(Header {
            name: Header::_CONTENT_TYPE.to_string(),
            value: "text/event-stream".to_string(),
        });
        response.headers.push(Header {
            name: Header::_CACHE_CONTROL.to_string(),
            value: "no-cache".to_string(),
        });
        response.headers.push(Header {
            name: "X-Accel-Buffering".to_string(),
            value: "no".to_string(),
        });
        response.stream_pipe = Some(Box::new(SseChannelReader::new(rx)));
        response
    }

    /// Advertise the `logging` capability (`"logging":{}`) in `initialize`'s
    /// response, so clients know they can call `logging/setLevel` and expect
    /// `notifications/message` log entries over the `GET /mcp` SSE stream.
    ///
    /// This only changes what's *advertised* — [`Self::log`] pushes log
    /// entries regardless of whether this was called, same as [`Self::notify`]
    /// works unconditionally. A spec-honest client simply wouldn't send
    /// `logging/setLevel` in the first place without seeing the capability.
    ///
    /// ```rust
    /// use rust_web_server::mcp::McpServer;
    ///
    /// let server = McpServer::new("my-server", "1.0").logging_enabled();
    /// ```
    pub fn logging_enabled(mut self) -> Self {
        self.logging_enabled = true;
        self
    }

    /// Push a `notifications/message` log entry to every client connected to
    /// the `GET /mcp` SSE stream, if `level` is at or above the level most
    /// recently set via a client's `logging/setLevel` request (or
    /// [`LogLevel::Debug`] — i.e. every level — if none has been set yet).
    ///
    /// `data_json` must already be a valid JSON value (the spec allows any
    /// type here, not just an object — a plain string is fine) — it's
    /// spliced in verbatim, not escaped or re-serialized. `logger`, if
    /// given, identifies the log's source (e.g. a module or subsystem name)
    /// and is escaped automatically.
    ///
    /// ```rust
    /// use rust_web_server::mcp::{LogLevel, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0").logging_enabled();
    /// server.log(LogLevel::Warning, Some("database"), r#""connection pool exhausted""#);
    /// ```
    pub fn log(&self, level: LogLevel, logger: Option<&str>, data_json: &str) {
        if level < *self.min_log_level.lock().unwrap() {
            return;
        }
        let logger_field = match logger {
            Some(l) => format!(r#","logger":"{}""#, json_escape(l)),
            None => String::new(),
        };
        let params = format!(r#"{{"level":"{}"{logger_field},"data":{data_json}}}"#, level.as_str());
        self.notify("notifications/message", Some(&params));
    }

    /// Handle `logging/setLevel`: store the requested minimum level so
    /// subsequent [`Self::log`] calls filter against it. Returns
    /// `INVALID_PARAMS` for a missing or unrecognized `params.level`.
    fn do_set_log_level(&self, body: &str) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let level_str = json_rpc::extract_str(&params, "level")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing level".to_string()))?;
        let level = LogLevel::parse(&level_str)
            .ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown log level: {level_str}")))?;
        *self.min_log_level.lock().unwrap() = level;
        Ok("{}".to_string())
    }

    // ── dynamic registration ──────────────────────────────────────────────────
    //
    // Unlike `.tool()`/`.resource()`/`.prompt()` (consuming builders, called
    // before the server starts serving requests), these take `&self` and can
    // be called at any time from any thread holding a clone of this
    // `McpServer` — e.g. after discovering a plugin, connecting to a
    // database, or reacting to a hot-reloaded config file. Every clone
    // shares the same underlying `Arc<RwLock<Vec<_>>>`, so a mutation made
    // through one clone is immediately visible to every other clone,
    // including the ones handling concurrent requests on other threads.
    //
    // Each registration/removal pushes the corresponding
    // `notifications/{tools,resources,prompts}/list_changed` event (no
    // params, per spec) to every `GET /mcp` SSE client via `.notify()`.

    /// Register a callable tool at runtime, exactly like [`Self::tool`] but
    /// without needing to own the server (and usable after it's already
    /// serving requests). Pushes `notifications/tools/list_changed`.
    ///
    /// ```rust
    /// use rust_web_server::mcp::{McpContent, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0");
    ///
    /// // Later, from any thread holding a clone of `server`:
    /// server.register_tool("refresh_cache", "Reload the in-memory cache", "{}", |_args| {
    ///     Ok(McpContent::text("cache refreshed"))
    /// });
    /// let existed = server.remove_tool("refresh_cache");
    /// assert!(existed);
    /// ```
    pub fn register_tool<F>(&self, name: &str, description: &str, input_schema: &str, handler: F)
    where
        F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.tools.write().unwrap().push(ToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: None,
            handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
        });
        self.notify("notifications/tools/list_changed", None);
    }

    /// Register an `async fn` tool handler at runtime, exactly like
    /// [`Self::async_tool`] but usable after the server is already serving
    /// requests. Requires the `http2` feature — see `async_tool`'s docs for
    /// the async-to-sync bridging details.
    #[cfg(feature = "http2")]
    pub fn register_async_tool<F, Fut>(&self, name: &str, description: &str, input_schema: &str, handler: F)
    where
        F: Fn(&str) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<McpContent, String>> + Send + 'static,
    {
        self.async_tools.write().unwrap().push(AsyncToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: None,
            handler: Arc::new(move |args: String| Box::pin(handler(&args))),
        });
        self.notify("notifications/tools/list_changed", None);
    }

    /// Remove a previously-registered tool by name — a sync tool registered
    /// via [`Self::tool`]/[`Self::register_tool`], or (with the `http2`
    /// feature) an async one via [`Self::async_tool`]/[`Self::register_async_tool`];
    /// both collections are checked. Returns `true` if a tool with that name
    /// existed in either and was removed. Pushes `notifications/tools/list_changed`
    /// only when something was actually removed.
    pub fn remove_tool(&self, name: &str) -> bool {
        #[cfg_attr(not(feature = "http2"), allow(unused_mut))]
        let mut removed = {
            let mut tools = self.tools.write().unwrap();
            let before = tools.len();
            tools.retain(|t| t.name != name);
            tools.len() != before
        };

        #[cfg(feature = "http2")]
        {
            let mut async_tools = self.async_tools.write().unwrap();
            let before = async_tools.len();
            async_tools.retain(|t| t.name != name);
            removed = removed || async_tools.len() != before;
        }

        if removed {
            self.notify("notifications/tools/list_changed", None);
        }
        removed
    }

    /// Register a readable resource at runtime, exactly like [`Self::resource`].
    /// Pushes `notifications/resources/list_changed`.
    pub fn register_resource<F>(&self, uri_template: &str, name: &str, description: &str, handler: F)
    where
        F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.resources.write().unwrap().push(ResourceDef {
            uri_template: uri_template.to_string(),
            name: name.to_string(),
            description: description.to_string(),
            handler: Arc::new(handler),
        });
        self.notify("notifications/resources/list_changed", None);
    }

    /// Remove a previously-registered resource by its exact `uri_template`
    /// (the same string passed to [`Self::register_resource`]/[`Self::resource`],
    /// not a concrete URI). Returns `true` if it existed. Pushes
    /// `notifications/resources/list_changed` only when something was removed.
    pub fn remove_resource(&self, uri_template: &str) -> bool {
        let removed = {
            let mut resources = self.resources.write().unwrap();
            let before = resources.len();
            resources.retain(|r| r.uri_template != uri_template);
            resources.len() != before
        };
        if removed {
            self.notify("notifications/resources/list_changed", None);
        }
        removed
    }

    /// Register a prompt template at runtime, exactly like [`Self::prompt`]
    /// (no argument definitions — use [`Self::remove_prompt`] +
    /// [`Self::register_prompt`] if you need to change a prompt's arguments
    /// later; there is no dynamic equivalent of [`Self::prompt_with_args`]).
    /// Pushes `notifications/prompts/list_changed`.
    pub fn register_prompt<F>(&self, name: &str, description: &str, handler: F)
    where
        F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
    {
        self.prompts.write().unwrap().push(PromptDef {
            name: name.to_string(),
            description: description.to_string(),
            arguments: vec![],
            handler: Arc::new(handler),
        });
        self.notify("notifications/prompts/list_changed", None);
    }

    /// Remove a previously-registered prompt by name. Returns `true` if it
    /// existed. Pushes `notifications/prompts/list_changed` only when
    /// something was removed.
    pub fn remove_prompt(&self, name: &str) -> bool {
        let removed = {
            let mut prompts = self.prompts.write().unwrap();
            let before = prompts.len();
            prompts.retain(|p| p.name != name);
            prompts.len() != before
        };
        if removed {
            self.notify("notifications/prompts/list_changed", None);
        }
        removed
    }

    /// Require a bearer token on every request to the MCP endpoint.
    ///
    /// The client must send `Authorization: Bearer <token>`. Requests with a
    /// missing or wrong token receive `401 Unauthorized` before any JSON-RPC
    /// processing occurs.
    ///
    /// Store the token in an environment variable — never hard-code it:
    ///
    /// ```rust,no_run
    /// use rust_web_server::app::App;
    /// use rust_web_server::core::New;
    ///
    /// let app = App::new()
    ///     .mcp("my-server", "1.0")
    ///     .require_bearer(std::env::var("MCP_TOKEN").expect("MCP_TOKEN not set"));
    /// ```
    ///
    /// Claude Desktop config:
    /// ```json
    /// { "mcpServers": { "my-server": {
    ///     "url": "http://localhost:7878/mcp",
    ///     "headers": { "Authorization": "Bearer <token>" }
    /// }}}
    /// ```
    ///
    /// An empty token refuses every request (fails closed) rather than
    /// accepting an empty `Authorization: Bearer` (since 17.113.1). Tokens
    /// are compared in constant time.
    pub fn require_bearer(mut self, token: impl Into<String>) -> Self {
        self.auth_token = Some(token.into());
        self
    }

    /// Browser origins allowed to call the MCP endpoint, e.g.
    /// `["https://inspector.example.com"]`. A request carrying an `Origin`
    /// header that isn't listed gets `403` (the MCP transport spec requires
    /// validating `Origin`: without it a web page — or a DNS-rebinding
    /// attack — can drive a local MCP server from the visitor's browser).
    /// Requests without `Origin` (CLI and agent clients) are unaffected.
    /// Empty by default: no browser origin is allowed (since 17.113.1).
    pub fn allowed_origins<I, S>(mut self, origins: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.allowed_origins = origins.into_iter().map(Into::into).collect();
        self
    }

    /// Require an OAuth 2.0 / OIDC bearer JWT — verified against a live
    /// JWKS endpoint — on every request to the MCP endpoint (MCP
    /// 2025-03-26's OAuth 2.0 authorization requirement), instead of one
    /// static shared secret. Requires the `sso` feature.
    ///
    /// Reuses [`crate::sso::jwks::JwksCache`] directly — the "leverage
    /// point" this feature was scoped around — rather than reimplementing
    /// JWT/JWKS verification. `provider` supplies the issuer, JWKS URI, and
    /// (for [`Self::execute`]'s `GET /.well-known/oauth-authorization-server`
    /// response) the authorization/token endpoints; use a preset
    /// (`OidcProvider::google()`, etc.) or `OidcProvider::discover(issuer)?`.
    ///
    /// A request with a missing, malformed, or invalid-signature bearer
    /// token gets `401` with `WWW-Authenticate: Bearer` before any JSON-RPC
    /// processing occurs — matching [`Self::require_bearer`]'s behavior. On
    /// success, the verified claims are available to `.tool_with_context()`
    /// handlers as JSON via [`McpContext::auth_claims`].
    ///
    /// If both this and [`Self::require_bearer`] are configured, OAuth
    /// verification takes precedence and the static token is never checked
    /// — configuring both is not a supported combination, just an
    /// unambiguous tie-break for an unusual configuration.
    ///
    /// ```rust,no_run
    /// use rust_web_server::app::App;
    /// use rust_web_server::core::New;
    /// use rust_web_server::sso::OidcProvider;
    ///
    /// let app = App::new()
    ///     .mcp("my-server", "1.0")
    ///     .require_oauth(OidcProvider::google(), "my-mcp-client-id");
    /// ```
    #[cfg(feature = "sso")]
    pub fn require_oauth(mut self, provider: crate::sso::OidcProvider, audience: impl Into<String>) -> Self {
        self.oauth = Some(Arc::new(OAuthConfig {
            jwks: crate::sso::jwks::JwksCache::new(&provider.jwks_uri),
            provider,
            audience: audience.into(),
        }));
        self
    }

    /// Wrap an existing [`Application`] so that non-MCP requests are forwarded
    /// to it instead of the built-in [`App`].
    ///
    /// Use this when your existing server has custom routes, state, or
    /// middleware that you want to keep alongside the MCP endpoint:
    ///
    /// ```rust,no_run
    /// use rust_web_server::app::App;
    /// use rust_web_server::mcp::{McpServer, McpContent};
    /// use rust_web_server::response::{Response, STATUS_CODE_REASON_PHRASE};
    /// use rust_web_server::test_client::TestClient;
    ///
    /// let existing_app = App::with_state(42u32)
    ///     .get("/api/hello", |_req, _params, _conn, _state| {
    ///         Response::get_response(&STATUS_CODE_REASON_PHRASE.n200_ok, None, None)
    ///     });
    ///
    /// let server = McpServer::new("my-app", "1.0")
    ///     .tool("ping", "Ping", "{}", |_| Ok(McpContent::text("pong")))
    ///     .wrap(existing_app);
    ///
    /// // Both /mcp and /api/hello are now handled by the same server.
    /// let client = TestClient::new(server);
    /// ```
    pub fn wrap(mut self, app: impl Application + Send + Sync + 'static) -> Self {
        self.fallback = Some(Arc::new(app));
        self
    }

    /// Override the HTTP path for the MCP endpoint (default `"/mcp"`).
    pub fn at(mut self, path: impl Into<String>) -> Self {
        self.path = path.into();
        self
    }

    /// Register a callable tool.
    ///
    /// - `name` — tool identifier (snake_case recommended)
    /// - `description` — human-readable description shown to the AI
    /// - `input_schema` — JSON Schema object for the tool's arguments
    /// - `handler` — closure receiving the raw `arguments` JSON string
    ///
    /// The handler returns [`McpContent`] on success or an error string.  An
    /// error is returned to the client as `isError: true` (not a protocol error).
    ///
    /// Use [`Self::tool_with_context`] instead if the handler needs the
    /// caller's identity, session, or headers.
    pub fn tool<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
    where
        F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.tools.write().unwrap().push(ToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: None,
            handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
        });
        self
    }

    /// Register an `async fn` tool handler — for a tool whose work is
    /// naturally async (an `AsyncClient` HTTP call, an async database
    /// query, awaiting another future) rather than forcing it through a
    /// blocking call inside a plain [`Self::tool`] handler.
    ///
    /// Requires the `http2` feature (tokio). `tools/call` bridges into the
    /// handler's future via [`crate::async_bridge::block_on_isolated`] —
    /// the same mechanism [`crate::proxy::H2ReverseProxy`] and
    /// [`crate::async_state::AsyncAppWithState`] already use to call async
    /// code from a synchronous [`Application::execute`] — rather than
    /// `tokio::task::block_in_place`, which panics outside a `multi_thread`
    /// runtime; `block_on_isolated` works under any runtime flavor,
    /// including `current_thread`.
    ///
    /// Like [`Self::tool`] (not [`Self::tool_with_context`]), the handler
    /// only receives `arguments` — there is no async equivalent of
    /// `tool_with_context` or `tool_annotated` yet.
    ///
    /// ```rust,no_run
    /// # #[cfg(feature = "http2")]
    /// # {
    /// use rust_web_server::mcp::{McpContent, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .async_tool("call_api", "Call an external API", "{}", |_args: &str| async move {
    ///         Ok(McpContent::text("response"))
    ///     });
    /// # }
    /// ```
    #[cfg(feature = "http2")]
    pub fn async_tool<F, Fut>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
    where
        F: Fn(&str) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<McpContent, String>> + Send + 'static,
    {
        self.async_tools.write().unwrap().push(AsyncToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: None,
            handler: Arc::new(move |args: String| Box::pin(handler(&args))),
        });
        self
    }

    /// Register a callable tool with [`ToolAnnotations`] — behavioral hints
    /// (read-only, destructive, idempotent, open-world) that MCP clients use
    /// to decide whether to warn or confirm before calling it. Otherwise
    /// identical to [`Self::tool`] — the handler still only receives
    /// `arguments`, not [`McpContext`] (there is currently no single builder
    /// combining annotations with per-request context; call [`Self::tool_with_context`]
    /// instead if you need context and don't need annotations).
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::{McpContent, McpServer, ToolAnnotations};
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_annotated(
    ///         "delete_file",
    ///         "Delete a file from disk",
    ///         r#"{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}"#,
    ///         ToolAnnotations {
    ///             destructive_hint: Some(true),
    ///             read_only_hint: Some(false),
    ///             idempotent_hint: Some(true), // deleting twice = deleting once
    ///             ..Default::default()
    ///         },
    ///         |_args| Ok(McpContent::text("deleted")),
    ///     );
    /// ```
    pub fn tool_annotated<F>(
        self,
        name: &str,
        description: &str,
        input_schema: &str,
        annotations: ToolAnnotations,
        handler: F,
    ) -> Self
    where
        F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.tools.write().unwrap().push(ToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: Some(annotations),
            handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
        });
        self
    }

    /// Register a callable tool whose handler also receives [`McpContext`] —
    /// caller identity/session info derived from this request's headers and
    /// whatever `clientInfo` this session sent at `initialize` time.
    ///
    /// Same `name`/`description`/`input_schema` semantics as [`Self::tool`];
    /// the only difference is the handler's first parameter.
    ///
    /// ```rust,no_run
    /// use rust_web_server::mcp::{McpContent, McpServer};
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .tool_with_context(
    ///         "whoami",
    ///         "Report the caller's client info",
    ///         "{}",
    ///         |ctx, _args| {
    ///             let name = ctx.client_name.as_deref().unwrap_or("unknown client");
    ///             Ok(McpContent::text(format!("Called by {name}")))
    ///         },
    ///     );
    /// ```
    pub fn tool_with_context<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
    where
        F: Fn(McpContext, &str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.tools.write().unwrap().push(ToolDef {
            name: name.to_string(),
            description: description.to_string(),
            input_schema: input_schema.to_string(),
            annotations: None,
            handler: Arc::new(handler),
        });
        self
    }

    /// Register a readable resource.
    ///
    /// `uri_template` uses `{param}` placeholders, e.g. `"user://{id}"`.
    /// The handler receives the full concrete URI string.
    pub fn resource<F>(self, uri_template: &str, name: &str, description: &str, handler: F) -> Self
    where
        F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
    {
        self.resources.write().unwrap().push(ResourceDef {
            uri_template: uri_template.to_string(),
            name: name.to_string(),
            description: description.to_string(),
            handler: Arc::new(handler),
        });
        self
    }

    /// Register a prompt template.
    ///
    /// The handler receives the raw `arguments` JSON string and returns a
    /// list of [`PromptMessage`] values.
    pub fn prompt<F>(self, name: &str, description: &str, handler: F) -> Self
    where
        F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
    {
        self.prompts.write().unwrap().push(PromptDef {
            name: name.to_string(),
            description: description.to_string(),
            arguments: vec![],
            handler: Arc::new(handler),
        });
        self
    }

    /// Register a prompt template with explicit argument definitions.
    pub fn prompt_with_args<F>(
        self,
        name: &str,
        description: &str,
        args: Vec<PromptArgDef>,
        handler: F,
    ) -> Self
    where
        F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
    {
        self.prompts.write().unwrap().push(PromptDef {
            name: name.to_string(),
            description: description.to_string(),
            arguments: args,
            handler: Arc::new(handler),
        });
        self
    }

    /// Register an argument-completion provider for one named argument of a
    /// tool or prompt, so clients like Cursor and VS Code can offer
    /// autocomplete while the user fills in that argument.
    ///
    /// `ref_type` is `"tool"` or `"prompt"` — matched against the incoming
    /// `completion/complete` request's `ref.type` (`"ref/tool"`/`"ref/prompt"`
    /// on the wire) with the `"ref/"` prefix stripped. `ref_name` is the
    /// tool or prompt name this applies to. The handler receives the
    /// argument's name and whatever partial value the user has typed so
    /// far, and returns candidate completion strings (or an error, mapped
    /// to a JSON-RPC `INVALID_PARAMS` response).
    ///
    /// `initialize` advertises the `completions` capability automatically
    /// once at least one `.completion()` has been registered — there's no
    /// separate opt-in flag to remember.
    ///
    /// ```rust
    /// use rust_web_server::mcp::McpServer;
    ///
    /// let server = McpServer::new("my-server", "1.0")
    ///     .completion("tool", "deploy", |arg_name, _partial| {
    ///         match arg_name {
    ///             "region" => Ok(vec!["us-east-1".to_string(), "eu-west-1".to_string()]),
    ///             _ => Ok(vec![]),
    ///         }
    ///     });
    /// ```
    pub fn completion<F>(self, ref_type: &str, ref_name: &str, handler: F) -> Self
    where
        F: Fn(&str, &str) -> Result<Vec<String>, String> + Send + Sync + 'static,
    {
        self.completions.write().unwrap().push(CompletionDef {
            ref_type: ref_type.to_string(),
            ref_name: ref_name.to_string(),
            handler: Arc::new(handler),
        });
        self
    }

    // ── request dispatch ──────────────────────────────────────────────────────

    /// Process a raw JSON-RPC body and return an HTTP response.
    ///
    /// Equivalent to [`Self::handle_request_with_context`] with an empty
    /// [`McpContext`] — tool handlers registered via
    /// [`Self::tool_with_context`] will see every field as `None`. Prefer
    /// calling through [`Application::execute`] (i.e. actually serving HTTP
    /// requests) when you need real per-request context; this method exists
    /// for calling the JSON-RPC layer directly, e.g. in tests.
    pub fn handle_request(&self, body: &str) -> Response {
        self.handle_request_with_context(body, McpContext::default())
    }

    /// Process a raw JSON-RPC body with an explicit [`McpContext`] and return
    /// an HTTP response. [`Self::execute`] calls this with a context built
    /// from the request's headers and this session's stored `clientInfo`;
    /// [`Self::handle_request`] calls this with an empty context.
    ///
    /// On a successful `initialize`, this mints a new session id (reusing
    /// [`crate::request_id::generate_request_id`]'s ID generator), records
    /// `params.clientInfo` under it, and returns the id in an
    /// `Mcp-Session-Id` response header — the client is expected to echo that
    /// header back on subsequent requests so later `tools/call`s in the same
    /// session can look their `clientInfo` back up.
    pub fn handle_request_with_context(&self, body: &str, ctx: McpContext) -> Response {
        let trimmed = body.trim_start();
        if trimmed.starts_with('[') {
            return self.handle_batch(trimmed, ctx);
        }

        let id = json_rpc::extract_id(body);

        let method = match json_rpc::extract_str(body, "method") {
            Some(m) => m,
            None => {
                // No `method` — either malformed, or a JSON-RPC *response*
                // to a sampling/createMessage request this server sent
                // earlier (responses never carry `method`).
                if let Some(id) = &id {
                    if self.try_deliver_sampling_response(id, body) {
                        return no_content();
                    }
                }
                return rpc_error(None, json_rpc::INVALID_REQUEST, "Missing method");
            }
        };

        if method == "notifications/cancelled" {
            self.handle_cancellation(body);
            return no_content();
        }

        if method == "notifications/roots/list_changed" {
            self.invalidate_roots_cache(&ctx.session_id);
            return no_content();
        }

        // Notifications have no `id` — acknowledge with 202 and no body.
        if method == "notifications/initialized" || (id.is_none() && method != "ping") {
            return no_content();
        }

        let result = self.dispatch_with_cancellation(&method, body, ctx, &id);
        let id_str = id.as_deref().unwrap_or("null");
        let is_ok = result.is_ok();

        let mut response = json_response(&Self::format_result(id_str, &result));

        if method == "initialize" && is_ok {
            self.start_session(body, &mut response);
        }

        response
    }

    /// Process a JSON-RPC 2.0 batch request — a top-level JSON array of
    /// request objects sent in a single `POST /mcp` body, per the JSON-RPC
    /// batch spec that MCP inherits. Each element is dispatched exactly as
    /// [`Self::handle_request_with_context`] would dispatch it standalone;
    /// notifications (no `id`) contribute no entry to the response array,
    /// same as they'd get no response body outside a batch.
    ///
    /// An empty array (`[]`) is itself an invalid request per the JSON-RPC
    /// spec, so it gets one `Invalid Request` error object back rather than
    /// an empty array. A batch made up entirely of notifications produces no
    /// response body at all (`202 Accepted`), matching a single notification.
    ///
    /// If the batch contains a successful `initialize` call, the *first* one
    /// mints a session and attaches `Mcp-Session-Id` to the overall response,
    /// same as a standalone `initialize` would — sending more than one
    /// `initialize` in a batch is unusual and only the first is honored for
    /// session purposes, since one HTTP response can only carry one session id.
    fn handle_batch(&self, array_body: &str, ctx: McpContext) -> Response {
        let elements = json_rpc::split_array_elements(array_body);
        if elements.is_empty() {
            return rpc_error(None, json_rpc::INVALID_REQUEST, "Invalid Request");
        }

        let mut parts: Vec<String> = Vec::new();
        let mut session_init_body: Option<String> = None;

        for elem in &elements {
            let id = json_rpc::extract_id(elem);

            let method = match json_rpc::extract_str(elem, "method") {
                Some(m) => m,
                None => {
                    if let Some(id) = &id {
                        if self.try_deliver_sampling_response(id, elem) {
                            continue; // delivered; no response entry, same as any notification
                        }
                    }
                    parts.push(Self::format_result(
                        "null",
                        &Err((json_rpc::INVALID_REQUEST, "Missing method".to_string())),
                    ));
                    continue;
                }
            };

            if method == "notifications/cancelled" {
                self.handle_cancellation(elem);
                continue;
            }

            if method == "notifications/roots/list_changed" {
                self.invalidate_roots_cache(&ctx.session_id);
                continue;
            }

            if method == "notifications/initialized" || (id.is_none() && method != "ping") {
                continue;
            }

            let result = self.dispatch_with_cancellation(&method, elem, ctx.clone(), &id);
            let id_str = id.as_deref().unwrap_or("null");
            let is_ok = result.is_ok();

            if method == "initialize" && is_ok && session_init_body.is_none() {
                session_init_body = Some(elem.clone());
            }

            parts.push(Self::format_result(id_str, &result));
        }

        if parts.is_empty() {
            // Every element was a notification — no response body, same as a
            // single standalone notification.
            return no_content();
        }

        let mut response = json_response(&format!("[{}]", parts.join(",")));
        if let Some(init_body) = session_init_body {
            self.start_session(&init_body, &mut response);
        }
        response
    }

    /// Dispatch one already-parsed JSON-RPC `method` against `body` (the raw
    /// single-object message, whether it arrived standalone or as one element
    /// of a batch) and return the JSON-RPC `result` payload or an error.
    /// Shared by [`Self::handle_request_with_context`] and [`Self::handle_batch`]
    /// so the method table exists in exactly one place.
    fn dispatch(&self, method: &str, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
        match method {
            "initialize"     => self.do_initialize(body),
            "ping"           => Ok("{}".to_string()),
            "tools/list"     => self.do_tools_list(body),
            "tools/call"     => self.do_tools_call(body, ctx),
            "resources/list" => self.do_resources_list(body),
            "resources/read" => self.do_resources_read(body),
            "resources/subscribe"   => self.do_resource_subscribe(body, ctx.session_id),
            "resources/unsubscribe" => self.do_resource_unsubscribe(body, ctx.session_id),
            "prompts/list"   => self.do_prompts_list(body),
            "prompts/get"    => self.do_prompts_get(body),
            "logging/setLevel" => self.do_set_log_level(body),
            "completion/complete" => self.do_completion(body),
            _                => Err((json_rpc::METHOD_NOT_FOUND, format!("Unknown method: {method}"))),
        }
    }

    /// Wraps [`Self::dispatch`] for one request, additionally registering a
    /// cancellation flag for `tools/call` (keyed by `id`, guaranteed `Some`
    /// for `tools/call` — a notification-shaped `tools/call` with no `id`
    /// never reaches this far, see the notification checks in
    /// [`Self::handle_request_with_context`]/[`Self::handle_batch`]) so a
    /// later `notifications/cancelled` referencing this exact request can
    /// flip it. The entry is removed again once dispatch returns, whether
    /// the handler checked the flag or not — this map never accumulates
    /// stale entries the way `sessions`/`sse_clients` can.
    fn dispatch_with_cancellation(
        &self,
        method: &str,
        body: &str,
        ctx: McpContext,
        id: &Option<String>,
    ) -> Result<String, (i32, String)> {
        let Some(id_str) = (method == "tools/call").then_some(id.as_ref()).flatten() else {
            return self.dispatch(method, body, ctx);
        };

        let flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
        self.cancellations.lock().unwrap().insert(id_str.clone(), flag.clone());
        let ctx = McpContext { cancellation: Some(flag), ..ctx };
        let result = self.dispatch(method, body, ctx);
        self.cancellations.lock().unwrap().remove(id_str);
        result
    }

    /// Handle `notifications/cancelled`: look up `params.requestId` (the
    /// spec allows `string | number`, so this is matched as a raw JSON
    /// token, the same way [`McpContext::progress_token`] is) in
    /// `cancellations` and flip its flag if the request is still in flight.
    /// Silently does nothing for an unknown/already-finished request id —
    /// the target call may simply have already completed naturally, which
    /// isn't an error.
    fn handle_cancellation(&self, body: &str) {
        let Some(params) = json_rpc::extract_raw(body, "params") else { return };
        let Some(request_id) = json_rpc::extract_raw(&params, "requestId") else { return };
        if let Some(flag) = self.cancellations.lock().unwrap().get(&request_id) {
            flag.store(true, std::sync::atomic::Ordering::Relaxed);
        }
    }

    /// Handle `notifications/roots/list_changed`: clear this session's
    /// cached `roots/list` result (if any), so the next
    /// [`McpContext::list_roots`] call in this session does a fresh round
    /// trip instead of serving stale data. Correlated purely by this
    /// request's own `Mcp-Session-Id` (`session_id`) — unlike
    /// `notifications/cancelled`, this notification carries no params of
    /// its own to key off of; it's about "your cache for *this connection*
    /// is stale," not any other request.
    fn invalidate_roots_cache(&self, session_id: &Option<String>) {
        let Some(sid) = session_id else { return };
        if let Some(info) = self.sessions.lock().unwrap().get_mut(sid) {
            info.roots = None;
        }
    }

    /// If `body` is a JSON-RPC *response* (no `method`, by definition — see
    /// [`Self::handle_request_with_context`]) to a `sampling/createMessage`
    /// request this server sent and is still waiting on (`id` found in
    /// `pending_replies`), deliver it to the blocked [`McpContext::sample`]
    /// call and return `true`. Returns `false` for any `id` this server
    /// isn't tracking — the caller still needs to treat that as a genuine
    /// malformed/unrecognized message.
    fn try_deliver_sampling_response(&self, id: &str, body: &str) -> bool {
        let sender = match self.pending_replies.lock().unwrap().remove(id) {
            Some(s) => s,
            None => return false,
        };
        let payload = match json_rpc::extract_raw(body, "result") {
            Some(result) => Ok(result),
            None => {
                let message = json_rpc::extract_raw(body, "error")
                    .and_then(|e| json_rpc::extract_str(&e, "message"))
                    .unwrap_or_else(|| "sampling/createMessage request failed".to_string());
                Err(message)
            }
        };
        let _ = sender.send(payload); // the waiter may have already timed out and dropped rx
        true
    }

    /// Render one JSON-RPC 2.0 response object — `{"jsonrpc":"2.0","result":...,"id":...}`
    /// or the `error` shape — from a dispatch result and its request's `id` (already
    /// rendered as a raw JSON token, e.g. `"1"`, `"\"abc\""`, or `"null"`).
    fn format_result(id_str: &str, result: &Result<String, (i32, String)>) -> String {
        match result {
            Ok(result_json) => format!(
                r#"{{"jsonrpc":"2.0","result":{result_json},"id":{id_str}}}"#
            ),
            Err((code, msg)) => {
                let escaped = json_escape(msg);
                format!(
                    r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
                )
            }
        }
    }

    /// Mint a new session id, record `body`'s `params.clientInfo` under it
    /// (logging the caller's identity), and attach the id to `response` as
    /// an `Mcp-Session-Id` header. Called once, from
    /// [`Self::handle_request_with_context`], right after a successful
    /// `initialize`.
    fn start_session(&self, body: &str, response: &mut Response) {
        let params = json_rpc::extract_raw(body, "params");
        let client_info = params.as_deref().and_then(|p| json_rpc::extract_raw(p, "clientInfo"));
        let (name, version) = match &client_info {
            Some(info) => (
                json_rpc::extract_str(info, "name"),
                json_rpc::extract_str(info, "version"),
            ),
            None => (None, None),
        };
        let capabilities = params.as_deref().and_then(|p| json_rpc::extract_raw(p, "capabilities"));
        let supports_sampling = capabilities.as_deref()
            .and_then(|c| json_rpc::extract_raw(c, "sampling"))
            .is_some();
        let supports_roots = capabilities.as_deref()
            .and_then(|c| json_rpc::extract_raw(c, "roots"))
            .is_some();

        eprintln!(
            "[mcp] initialize from client {} v{}",
            name.as_deref().unwrap_or("unknown"),
            version.as_deref().unwrap_or("unknown"),
        );

        let session_id = crate::request_id::generate_request_id();
        self.sessions
            .lock()
            .unwrap()
            .insert(session_id.clone(), StoredClientInfo { name, version, supports_sampling, supports_roots, roots: None });

        response.headers.push(Header {
            name: "Mcp-Session-Id".to_string(),
            value: session_id,
        });
    }

    // ── method handlers ───────────────────────────────────────────────────────

    /// Handle `initialize`. Per spec, the server must inspect the client's
    /// requested `protocolVersion` and respond with the version it actually
    /// supports — allowing the client to abort the session if incompatible —
    /// rather than blindly echoing `PROTOCOL_VERSION` regardless of what was
    /// asked for.
    ///
    /// This server implements exactly one protocol version, so "negotiation"
    /// here means: if the client asked for that same version, confirm it;
    /// otherwise tell the client the version we actually speak (older *or*
    /// newer than what was requested), which is always the lower of the two
    /// — version strings are `YYYY-MM-DD` dates, so a plain string comparison
    /// already orders them correctly with no date parsing needed.
    ///
    /// `clientInfo` is *not* handled here — [`Self::handle_request_with_context`]
    /// extracts and stores it (under a freshly minted session id) after this
    /// returns, so it's only ever parsed out of `body` once per call.
    fn do_initialize(&self, body: &str) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params");

        let client_version = params.as_deref().and_then(|p| json_rpc::extract_str(p, "protocolVersion"));

        let negotiated_version: &str = match client_version.as_deref() {
            Some(v) if v < PROTOCOL_VERSION => v,
            _ => PROTOCOL_VERSION,
        };

        let logging_cap = if self.logging_enabled { r#","logging":{}"# } else { "" };
        // completions is advertised iff at least one .completion() has been registered —
        // no separate opt-in flag needed, unlike logging: if nothing was registered,
        // completion/complete would just return empty results for everything anyway.
        let completions_cap = if self.completions.read().unwrap().is_empty() { "" } else { r#","completions":{}"# };
        // listChanged is always true: register_tool/remove_tool (and the resource/prompt
        // equivalents) are always available, unlike logging which is opt-in via
        // .logging_enabled(). resources.subscribe is also always true now that
        // resources/subscribe and resources/unsubscribe are implemented (MCP_TODO.md's TODO-14).
        let caps = format!(
            r#"{{"tools":{{"listChanged":true}},"resources":{{"subscribe":true,"listChanged":true}},"prompts":{{"listChanged":true}}{logging_cap}{completions_cap}}}"#
        );
        Ok(format!(
            r#"{{"protocolVersion":"{}","capabilities":{caps},"serverInfo":{{"name":"{}","version":"{}"}}}}"#,
            json_escape(negotiated_version),
            json_escape(&self.server_name),
            json_escape(&self.server_version),
        ))
    }

    fn do_tools_list(&self, body: &str) -> Result<String, (i32, String)> {
        #[cfg_attr(not(feature = "http2"), allow(unused_mut))]
        let mut items: Vec<String> = self.tools.read().unwrap().iter()
            .map(|t| Self::render_tool_list_entry(&t.name, &t.description, &t.input_schema, t.annotations))
            .collect();

        #[cfg(feature = "http2")]
        items.extend(self.async_tools.read().unwrap().iter()
            .map(|t| Self::render_tool_list_entry(&t.name, &t.description, &t.input_schema, t.annotations)));

        let (page, next_cursor) = self.paginate(&items, body)?;
        Ok(format!(r#"{{"tools":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
    }

    /// Render one `tools/list` entry — shared by the sync (`tools`) and
    /// (with the `http2` feature) async (`async_tools`) collections, which
    /// are otherwise identical from the client's point of view.
    fn render_tool_list_entry(name: &str, description: &str, input_schema: &str, annotations: Option<ToolAnnotations>) -> String {
        let annotations_field = match annotations {
            Some(a) => format!(r#","annotations":{}"#, a.to_json()),
            None => String::new(),
        };
        format!(
            r#"{{"name":"{}","description":"{}","inputSchema":{}{}}}"#,
            json_escape(name), json_escape(description), input_schema, annotations_field,
        )
    }

    fn do_tools_call(&self, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let name = json_rpc::extract_str(&params, "name")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing tool name".to_string()))?;
        let args = json_rpc::extract_raw(&params, "arguments")
            .unwrap_or_else(|| "{}".to_string());

        // `_meta.progressToken` (string or number, per spec) — stored raw so
        // `McpContext::report_progress` can splice it back verbatim.
        let progress_token = json_rpc::extract_raw(&params, "_meta")
            .and_then(|meta| json_rpc::extract_raw(&meta, "progressToken"));
        let ctx = McpContext { progress_token, ..ctx };

        let sync_handler = {
            let tools = self.tools.read().unwrap();
            tools.iter().find(|t| t.name == name).map(|t| t.handler.clone())
        };
        if let Some(handler) = sync_handler {
            return Ok(Self::format_tool_result(handler(ctx, &args)));
        }

        #[cfg(feature = "http2")]
        {
            let async_handler = {
                let async_tools = self.async_tools.read().unwrap();
                async_tools.iter().find(|t| t.name == name).map(|t| t.handler.clone())
            };
            if let Some(handler) = async_handler {
                let args_owned = args.clone();
                let result = crate::async_bridge::block_on_isolated(move || handler(args_owned));
                return Ok(Self::format_tool_result(result));
            }
        }

        Err((json_rpc::INVALID_PARAMS, format!("Unknown tool: {name}")))
    }

    /// Render a tool handler's result as the spec's `tools/call` shape —
    /// `{"content":[...],"isError":false}` on success, or `isError:true`
    /// with the error message as a text content block. Shared by the sync
    /// and (with `http2`) async `tools/call` dispatch paths.
    fn format_tool_result(result: Result<McpContent, String>) -> String {
        match result {
            Ok(c) => format!(
                r#"{{"content":[{}],"isError":false}}"#,
                c.to_content_json(),
            ),
            Err(e) => {
                let escaped = json_escape(&e);
                format!(
                    r#"{{"content":[{{"type":"text","text":"{escaped}"}}],"isError":true}}"#
                )
            }
        }
    }

    fn do_resources_list(&self, body: &str) -> Result<String, (i32, String)> {
        let items: Vec<String> = self.resources.read().unwrap().iter().map(|r| {
            format!(
                r#"{{"uri":"{}","name":"{}","description":"{}","mimeType":"text/plain"}}"#,
                json_escape(&r.uri_template),
                json_escape(&r.name),
                json_escape(&r.description),
            )
        }).collect();
        let (page, next_cursor) = self.paginate(&items, body)?;
        Ok(format!(r#"{{"resources":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
    }

    fn do_resources_read(&self, body: &str) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let uri = json_rpc::extract_str(&params, "uri")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;

        let handler = {
            let resources = self.resources.read().unwrap();
            resources.iter().find(|r| uri_matches(&r.uri_template, &uri)).map(|r| r.handler.clone())
        }.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Resource not found: {uri}")))?;

        match handler(&uri) {
            Ok(c) => {
                let text_esc = json_escape(&c.text);
                let uri_esc  = json_escape(&uri);
                Ok(format!(
                    r#"{{"contents":[{{"uri":"{uri_esc}","mimeType":"{}","text":"{text_esc}"}}]}}"#,
                    c.mime(),
                ))
            }
            Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
        }
    }

    /// Handle `resources/subscribe`: record that `session_id` (this
    /// request's `Mcp-Session-Id`) wants `notifications/resources/updated`
    /// pushed whenever [`Self::notify_resource_updated`] is called for
    /// `params.uri`. Requires a session id — without one there's no way to
    /// later match this subscription to a specific `GET /mcp` SSE
    /// connection, so a client that never sent `Mcp-Session-Id` gets
    /// `INVALID_PARAMS` rather than a subscription that can never fire.
    fn do_resource_subscribe(&self, body: &str, session_id: Option<String>) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let uri = json_rpc::extract_str(&params, "uri")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
        let session_id = session_id.ok_or((
            json_rpc::INVALID_PARAMS,
            "resources/subscribe requires an Mcp-Session-Id header (call initialize first)".to_string(),
        ))?;

        let mut subs = self.subscriptions.lock().unwrap();
        let subscribers = subs.entry(uri).or_default();
        if !subscribers.contains(&session_id) {
            subscribers.push(session_id);
        }
        Ok("{}".to_string())
    }

    /// Handle `resources/unsubscribe`: the inverse of
    /// [`Self::do_resource_subscribe`]. Removing the last subscriber for a
    /// URI drops the URI's entry entirely rather than leaving an empty
    /// `Vec` behind.
    fn do_resource_unsubscribe(&self, body: &str, session_id: Option<String>) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let uri = json_rpc::extract_str(&params, "uri")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
        let session_id = session_id.ok_or((
            json_rpc::INVALID_PARAMS,
            "resources/unsubscribe requires an Mcp-Session-Id header".to_string(),
        ))?;

        let mut subs = self.subscriptions.lock().unwrap();
        if let Some(subscribers) = subs.get_mut(&uri) {
            subscribers.retain(|s| s != &session_id);
            if subscribers.is_empty() {
                subs.remove(&uri);
            }
        }
        Ok("{}".to_string())
    }

    fn do_prompts_list(&self, body: &str) -> Result<String, (i32, String)> {
        let items: Vec<String> = self.prompts.read().unwrap().iter().map(|p| {
            let arg_defs: Vec<String> = p.arguments.iter().map(|a| {
                format!(
                    r#"{{"name":"{}","description":"{}","required":{}}}"#,
                    json_escape(&a.name),
                    json_escape(&a.description),
                    a.required,
                )
            }).collect();
            format!(
                r#"{{"name":"{}","description":"{}","arguments":[{}]}}"#,
                json_escape(&p.name),
                json_escape(&p.description),
                arg_defs.join(","),
            )
        }).collect();
        let (page, next_cursor) = self.paginate(&items, body)?;
        Ok(format!(r#"{{"prompts":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
    }

    /// Slice `items` (already-rendered JSON object strings for one
    /// `*/list` response) according to [`Self::page_size`] and this
    /// request's `params.cursor`, returning the page and — if more items
    /// remain — the opaque `nextCursor` to embed in the response.
    ///
    /// Without a configured `page_size`, always returns every item and no
    /// cursor, i.e. pagination is fully opt-in.
    fn paginate<'a>(&self, items: &'a [String], body: &str) -> Result<(&'a [String], Option<String>), (i32, String)> {
        let page_size = match self.page_size {
            Some(n) => n,
            None => return Ok((items, None)),
        };

        let cursor = json_rpc::extract_raw(body, "params")
            .and_then(|p| json_rpc::extract_str(&p, "cursor"));

        let offset = match cursor {
            Some(c) => decode_cursor(&c)
                .ok_or((json_rpc::INVALID_PARAMS, "Invalid cursor".to_string()))?,
            None => 0,
        };

        if offset >= items.len() {
            return Ok((&[], None));
        }

        let end = (offset + page_size).min(items.len());
        let next_cursor = if end < items.len() { Some(encode_cursor(end)) } else { None };
        Ok((&items[offset..end], next_cursor))
    }

    fn do_prompts_get(&self, body: &str) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let name = json_rpc::extract_str(&params, "name")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing prompt name".to_string()))?;
        let args = json_rpc::extract_raw(&params, "arguments")
            .unwrap_or_else(|| "{}".to_string());

        let (description, handler) = {
            let prompts = self.prompts.read().unwrap();
            prompts.iter().find(|p| p.name == name).map(|p| (p.description.clone(), p.handler.clone()))
        }.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown prompt: {name}")))?;

        match handler(&args) {
            Ok(msgs) => {
                let msg_jsons: Vec<String> = msgs.iter().map(|m| m.to_json()).collect();
                Ok(format!(
                    r#"{{"description":"{}","messages":[{}]}}"#,
                    json_escape(&description),
                    msg_jsons.join(","),
                ))
            }
            Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
        }
    }

    /// Handle `completion/complete`: look up the registered
    /// [`CompletionDef`] matching `params.ref.type`/`params.ref.name`, call
    /// its handler with `params.argument.name`/`params.argument.value`, and
    /// render the spec's `{"completion":{"values":[...],"hasMore":...,
    /// "total":...}}` shape. No registered provider for the given ref/name
    /// (or an unrecognized `ref.type` not stripped of `"ref/"`) returns an
    /// empty `values` list rather than an error — matching the spec's own
    /// framing of completion as a best-effort hint, not a required capability
    /// per tool/prompt.
    fn do_completion(&self, body: &str) -> Result<String, (i32, String)> {
        let params = json_rpc::extract_raw(body, "params")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
        let reference = json_rpc::extract_raw(&params, "ref")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing ref".to_string()))?;
        let ref_type_raw = json_rpc::extract_str(&reference, "type")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing ref.type".to_string()))?;
        let ref_type = ref_type_raw.strip_prefix("ref/").unwrap_or(&ref_type_raw);
        let ref_name = json_rpc::extract_str(&reference, "name")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing ref.name".to_string()))?;

        let argument = json_rpc::extract_raw(&params, "argument")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing argument".to_string()))?;
        let arg_name = json_rpc::extract_str(&argument, "name")
            .ok_or((json_rpc::INVALID_PARAMS, "Missing argument.name".to_string()))?;
        let partial = json_rpc::extract_str(&argument, "value").unwrap_or_default();

        let handler = {
            let completions = self.completions.read().unwrap();
            completions.iter()
                .find(|c| c.ref_type == ref_type && c.ref_name == ref_name)
                .map(|c| c.handler.clone())
        };

        let values = match handler {
            Some(h) => h(&arg_name, &partial).map_err(|e| (json_rpc::INVALID_PARAMS, e))?,
            None => vec![],
        };

        let total = values.len();
        let has_more = total > MAX_COMPLETION_VALUES;
        let page = if has_more { &values[..MAX_COMPLETION_VALUES] } else { &values[..] };
        let values_json: Vec<String> = page.iter().map(|v| format!(r#""{}""#, json_escape(v))).collect();

        Ok(format!(
            r#"{{"completion":{{"values":[{}],"hasMore":{has_more},"total":{total}}}}}"#,
            values_json.join(","),
        ))
    }

    /// Build the [`McpContext`] for an incoming request: the `Mcp-Session-Id`
    /// header, if present, plus whatever `clientInfo` was recorded for that
    /// session at `initialize` time (if this session is recognized).
    fn context_for(&self, request: &Request) -> McpContext {
        let session_id = request
            .get_header("Mcp-Session-Id".to_string())
            .map(|h| h.value.clone());

        let (client_name, client_version, sampling_supported, roots_supported) = match &session_id {
            Some(sid) => match self.sessions.lock().unwrap().get(sid) {
                Some(info) => (info.name.clone(), info.version.clone(), info.supports_sampling, info.supports_roots),
                None => (None, None, false, false),
            },
            None => (None, None, false, false),
        };

        McpContext {
            client_name,
            client_version,
            session_id,
            auth_claims: None,
            progress_token: None,
            sse_clients: Some(self.sse_clients.clone()),
            cancellation: None,
            sampling_supported,
            pending_replies: Some(self.pending_replies.clone()),
            roots_supported,
            sessions: Some(self.sessions.clone()),
        }
    }
}

/// Render one JSON-RPC 2.0 notification (no `id` — fire-and-forget, per
/// spec) as an SSE `data:`-ready message body. Shared by [`McpServer::notify`]
/// and [`McpContext::report_progress`].
fn render_notification(method: &str, params_json: Option<&str>) -> String {
    let params_field = match params_json {
        Some(p) => format!(r#","params":{p}"#),
        None => String::new(),
    };
    format!(r#"{{"jsonrpc":"2.0","method":"{}"{}}}"#, json_escape(method), params_field)
}

/// Send a raw pre-built JSON-RPC message to every client in `clients`,
/// pruning any whose channel is full or disconnected. Shared by
/// [`McpServer::notify`] and [`McpContext::report_progress`] — the latter
/// only has a clone of the broadcast list, not a whole `McpServer`.
fn broadcast_sse_to(clients: &Arc<Mutex<Vec<SseClient>>>, json: &str) {
    let frame = format!("data: {json}\n\n").into_bytes();
    let mut clients = clients.lock().unwrap();
    clients.retain(|c| c.sender.try_send(frame.clone()).is_ok());
}

/// Send a raw pre-built JSON-RPC message only to clients whose `session_id`
/// is in `session_ids`, pruning any of *those* whose channel is full or
/// disconnected. A client with no `session_id`, or one not in
/// `session_ids`, is left untouched (kept, not sent to) — this is a
/// targeted send for [`McpServer::notify_resource_updated`], not a
/// broadcast like [`broadcast_sse_to`].
fn send_sse_to_sessions(clients: &Arc<Mutex<Vec<SseClient>>>, session_ids: &[String], json: &str) {
    let frame = format!("data: {json}\n\n").into_bytes();
    let mut clients = clients.lock().unwrap();
    clients.retain(|c| {
        let is_target = c.session_id.as_deref().is_some_and(|sid| session_ids.iter().any(|s| s == sid));
        if is_target {
            c.sender.try_send(frame.clone()).is_ok()
        } else {
            true
        }
    });
}

// ── SSE channel reader ────────────────────────────────────────────────────────

/// Adapts an `mpsc::Receiver<Vec<u8>>` of pre-framed SSE bytes into a
/// blocking [`std::io::Read`], so `Server::pipe_stream` (already written for
/// proxy passthrough streaming) can drive a `GET /mcp` SSE connection with no
/// changes to the server's write loop.
///
/// Blocks in [`Self::read`] until either a frame arrives, the sender side is
/// dropped (all `McpServer` clones gone — EOF, closing the connection), or
/// [`SSE_KEEPALIVE_INTERVAL`] elapses with nothing to send (writes a `:
/// keep-alive` comment instead, both to satisfy proxies that time out
/// silent connections and to surface a dead peer on the next write attempt).
struct SseChannelReader {
    rx: mpsc::Receiver<Vec<u8>>,
    leftover: Vec<u8>,
}

impl SseChannelReader {
    fn new(rx: mpsc::Receiver<Vec<u8>>) -> Self {
        SseChannelReader { rx, leftover: Vec::new() }
    }
}

impl std::io::Read for SseChannelReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        if self.leftover.is_empty() {
            loop {
                match self.rx.recv_timeout(SSE_KEEPALIVE_INTERVAL) {
                    Ok(frame) => { self.leftover = frame; break; }
                    Err(mpsc::RecvTimeoutError::Timeout) => {
                        self.leftover = b": keep-alive\n\n".to_vec();
                        break;
                    }
                    Err(mpsc::RecvTimeoutError::Disconnected) => return Ok(0),
                }
            }
        }

        let n = self.leftover.len().min(buf.len());
        buf[..n].copy_from_slice(&self.leftover[..n]);
        self.leftover.drain(..n);
        Ok(n)
    }
}

// ── Application ───────────────────────────────────────────────────────────────

impl Application for McpServer {
    fn execute(&self, request: &Request, connection: &ConnectionInfo) -> Result<Response, String> {
        #[cfg(feature = "sso")]
        if request.request_uri == "/.well-known/oauth-authorization-server" && request.method == "GET" {
            if let Some(oauth) = &self.oauth {
                return Ok(oauth_metadata_document(oauth));
            }
        }

        if request.request_uri == self.path {
            // Browser requests (they carry `Origin`) only from listed origins.
            if let Some(origin) = request.headers.iter().find(|h| h.name.eq_ignore_ascii_case("origin")) {
                let origin = origin.value.trim();
                if !self.allowed_origins.iter().any(|o| o.trim_end_matches('/').eq_ignore_ascii_case(origin.trim_end_matches('/'))) {
                    return Ok(forbidden_origin());
                }
            }
            // JSON-RPC bodies only: a form or text/plain POST is what a web
            // page can send cross-site without a CORS preflight.
            if request.method == "POST" {
                let content_type = request.headers.iter()
                    .find(|h| h.name.eq_ignore_ascii_case("content-type"))
                    .map(|h| h.value.as_str())
                    .unwrap_or("");
                let media_type = content_type.split(';').next().unwrap_or("").trim();
                if !media_type.eq_ignore_ascii_case("application/json") {
                    return Ok(unsupported_media_type());
                }
            }

            #[cfg_attr(not(feature = "sso"), allow(unused_mut))]
            let mut auth_claims_json: Option<String> = None;

            // OAuth (sso feature) takes precedence over a static bearer
            // token if both are somehow configured — see require_oauth's
            // doc comment.
            #[cfg(feature = "sso")]
            if let Some(oauth) = &self.oauth {
                let provided = request.headers.iter()
                    .find(|h| h.name.eq_ignore_ascii_case("authorization"))
                    .map(|h| h.value.as_str())
                    .unwrap_or("");
                let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
                if bearer.is_empty() {
                    return Ok(unauthorized());
                }
                let opts = crate::sso::jwks::VerifyOptions {
                    audience: &oauth.audience,
                    issuer: &oauth.provider.issuer,
                    leeway_secs: 60,
                };
                match oauth.jwks.verify_jwt(bearer, &opts) {
                    Ok(claims) => auth_claims_json = serde_json::to_string(&claims).ok(),
                    Err(_) => return Ok(unauthorized()),
                }
            } else if let Some(expected) = &self.auth_token {
                // Check the static bearer token before processing any MCP request.
                let provided = request.headers.iter()
                    .find(|h| h.name.eq_ignore_ascii_case("authorization"))
                    .map(|h| h.value.as_str())
                    .unwrap_or("");
                let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
                if !bearer_matches(expected, bearer) {
                    return Ok(unauthorized());
                }
            }
            #[cfg(not(feature = "sso"))]
            if let Some(expected) = &self.auth_token {
                let provided = request.headers.iter()
                    .find(|h| h.name.eq_ignore_ascii_case("authorization"))
                    .map(|h| h.value.as_str())
                    .unwrap_or("");
                let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
                if !bearer_matches(expected, bearer) {
                    return Ok(unauthorized());
                }
            }

            return Ok(match request.method.as_str() {
                "POST" => {
                    let body = std::str::from_utf8(&request.body).unwrap_or("");
                    let mut ctx = self.context_for(request);
                    ctx.auth_claims = auth_claims_json;
                    self.handle_request_with_context(body, ctx)
                }
                "GET" => self.start_sse_stream(request),
                "OPTIONS" => {
                    // CORS preflight for browser-based MCP clients
                    let mut r = Response::new();
                    r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
                    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
                    r.headers.push(Header {
                        name: "Allow".to_string(),
                        value: "GET, POST, OPTIONS".to_string(),
                    });
                    r
                }
                _ => {
                    let mut r = Response::new();
                    r.status_code = *STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.status_code;
                    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.reason_phrase.to_string();
                    r.headers.push(Header {
                        name: "Allow".to_string(),
                        value: "GET, POST, OPTIONS".to_string(),
                    });
                    r.content_range_list = vec![Range::get_content_range(
                        b"MCP endpoint only accepts GET (SSE) or POST".to_vec(),
                        MimeType::TEXT_PLAIN.to_string(),
                    )];
                    r
                }
            });
        }

        // Not an MCP path — fall through to the wrapped app (or built-in App).
        match &self.fallback {
            Some(app) => app.execute(request, connection),
            None      => App::new().execute(request, connection),
        }
    }

    /// MCP calls run tool handlers (user code): never inline. Other requests
    /// go to the fallback app, which decides.
    fn runs_inline(&self, request: &Request, connection: &ConnectionInfo) -> bool {
        if request.request_uri == self.path || request.request_uri.starts_with("/.well-known/") {
            return false;
        }
        match &self.fallback {
            Some(app) => app.runs_inline(request, connection),
            None => App::is_inline_request(request),
        }
    }

    #[cfg(feature = "http2")]
    fn execute_async<'a>(&'a self, request: &'a Request, connection: &'a ConnectionInfo) -> Option<crate::application::AsyncResponse<'a>> {
        if request.request_uri == self.path || request.request_uri.starts_with("/.well-known/") {
            return None;
        }
        self.fallback.as_ref()?.execute_async(request, connection)
    }
}

// ── public helper ─────────────────────────────────────────────────────────────

/// Extract a string argument from a tool/prompt `arguments` JSON object.
///
/// ```rust
/// use rust_web_server::mcp::extract_arg;
/// assert_eq!(extract_arg(r#"{"text":"hello"}"#, "text").as_deref(), Some("hello"));
/// assert_eq!(extract_arg(r#"{}"#, "missing"), None);
/// ```
pub fn extract_arg(arguments: &str, name: &str) -> Option<String> {
    json_rpc::extract_str(arguments, name)
}

// ── internal helpers ──────────────────────────────────────────────────────────

fn json_response(body: &str) -> Response {
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
    r.content_range_list = vec![Range::get_content_range(
        body.as_bytes().to_vec(),
        MimeType::APPLICATION_JSON.to_string(),
    )];
    r
}

fn no_content() -> Response {
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n202_accepted.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n202_accepted.reason_phrase.to_string();
    r
}

/// A configured static token against the presented one: never matches an
/// empty token (fails closed), compared in constant time.
fn bearer_matches(expected: &str, provided: &str) -> bool {
    !expected.is_empty() && crate::proxy_config::constant_time_eq(expected.as_bytes(), provided.as_bytes())
}

/// `403` for a browser `Origin` not in [`McpServer::allowed_origins`].
fn forbidden_origin() -> Response {
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n403_forbidden.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n403_forbidden.reason_phrase.to_string();
    r.content_range_list = vec![Range::get_content_range(
        b"Origin not allowed".to_vec(),
        MimeType::TEXT_PLAIN.to_string(),
    )];
    r
}

/// `415` for a POST whose body isn't `application/json`.
fn unsupported_media_type() -> Response {
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n415_unsupported_media_type.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n415_unsupported_media_type.reason_phrase.to_string();
    r.content_range_list = vec![Range::get_content_range(
        b"MCP requests must be application/json".to_vec(),
        MimeType::TEXT_PLAIN.to_string(),
    )];
    r
}

fn unauthorized() -> Response {
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n401_unauthorized.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n401_unauthorized.reason_phrase.to_string();
    r.headers.push(Header {
        name: "WWW-Authenticate".to_string(),
        value: "Bearer".to_string(),
    });
    r.content_range_list = vec![Range::get_content_range(
        b"Unauthorized".to_vec(),
        MimeType::TEXT_PLAIN.to_string(),
    )];
    r
}

/// `GET /.well-known/oauth-authorization-server` — served whenever
/// [`McpServer::require_oauth`] is configured, per MCP 2025-03-26's OAuth
/// 2.0 authorization requirement. Not a full RFC 8414 document (this
/// server is a resource server, not the authorization server itself) —
/// just enough for a client to confirm the issuer, JWKS URI, and the
/// authorization/token endpoints already known from `provider`.
#[cfg(feature = "sso")]
fn oauth_metadata_document(oauth: &OAuthConfig) -> Response {
    let body = format!(
        r#"{{"issuer":"{}","authorization_endpoint":"{}","token_endpoint":"{}","jwks_uri":"{}","response_types_supported":["code"]}}"#,
        json_escape(&oauth.provider.issuer),
        json_escape(&oauth.provider.authorization_endpoint),
        json_escape(&oauth.provider.token_endpoint),
        json_escape(&oauth.provider.jwks_uri),
    );
    let mut r = Response::new();
    r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
    r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
    r.content_range_list = vec![Range::get_content_range(
        body.into_bytes(),
        MimeType::APPLICATION_JSON.to_string(),
    )];
    r
}

fn rpc_error(id: Option<&str>, code: i32, message: &str) -> Response {
    let id_str  = id.unwrap_or("null");
    let escaped = json_escape(message);
    json_response(&format!(
        r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
    ))
}

pub(crate) fn json_escape(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 4);
    for ch in s.chars() {
        match ch {
            '"'  => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => { let _ = std::fmt::Write::write_fmt(&mut out, format_args!("\\u{:04x}", c as u32)); }
            c    => out.push(c),
        }
    }
    out
}

// ── pagination cursors ─────────────────────────────────────────────────────────

/// Render `,"nextCursor":"..."` for a `*/list` response, or `""` if there's
/// no next page — spliced directly after the closing `]` of the items array.
fn next_cursor_json(next_cursor: &Option<String>) -> String {
    match next_cursor {
        Some(c) => format!(r#","nextCursor":"{}""#, json_escape(c)),
        None => String::new(),
    }
}

/// Encode a `tools/list`/`resources/list`/`prompts/list` offset as the
/// opaque `nextCursor`/`params.cursor` string the MCP spec expects — just
/// base64 of the decimal offset, e.g. `50` → `"NTA="`. Callers only ever
/// treat this as opaque; the encoding is a private implementation detail of
/// this module, not a client-facing contract.
fn encode_cursor(offset: usize) -> String {
    base64_encode(offset.to_string().as_bytes())
}

/// Decode a cursor produced by [`encode_cursor`]. Returns `None` for
/// anything that isn't valid base64 of a decimal `usize` — a malformed or
/// tampered cursor, not a crash.
fn decode_cursor(cursor: &str) -> Option<usize> {
    String::from_utf8(base64_decode(cursor)?).ok()?.parse().ok()
}

const BASE64_TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

fn base64_encode(data: &[u8]) -> String {
    let mut out = String::with_capacity((data.len() + 2) / 3 * 4);
    for chunk in data.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
        let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
        let n = (b0 << 16) | (b1 << 8) | b2;
        out.push(BASE64_TABLE[((n >> 18) & 0x3F) as usize] as char);
        out.push(BASE64_TABLE[((n >> 12) & 0x3F) as usize] as char);
        out.push(if chunk.len() > 1 { BASE64_TABLE[((n >> 6) & 0x3F) as usize] as char } else { '=' });
        out.push(if chunk.len() > 2 { BASE64_TABLE[(n & 0x3F) as usize] as char } else { '=' });
    }
    out
}

fn base64_decode(s: &str) -> Option<Vec<u8>> {
    fn sextet(c: u8) -> Option<u32> {
        match c {
            b'A'..=b'Z' => Some((c - b'A') as u32),
            b'a'..=b'z' => Some((c - b'a' + 26) as u32),
            b'0'..=b'9' => Some((c - b'0' + 52) as u32),
            b'+' => Some(62),
            b'/' => Some(63),
            _ => None,
        }
    }

    let trimmed = s.trim_end_matches('=');
    let bytes = trimmed.as_bytes();
    let mut out = Vec::with_capacity(bytes.len() * 3 / 4 + 3);
    for chunk in bytes.chunks(4) {
        if chunk.len() == 1 {
            return None; // not a valid base64 length
        }
        let vals: Vec<u32> = chunk.iter().map(|&b| sextet(b)).collect::<Option<Vec<_>>>()?;
        let n = vals.iter().enumerate().fold(0u32, |acc, (i, &v)| acc | (v << (18 - 6 * i)));
        out.push(((n >> 16) & 0xFF) as u8);
        if vals.len() > 2 { out.push(((n >> 8) & 0xFF) as u8); }
        if vals.len() > 3 { out.push((n & 0xFF) as u8); }
    }
    Some(out)
}

fn uri_matches(template: &str, uri: &str) -> bool {
    // Template `"user://{id}"` matches any URI starting with `"user://"`.
    match template.find('{') {
        Some(pos) => uri.starts_with(&template[..pos]),
        None      => template == uri,
    }
}