protocol_droid/application/server/

builder.rs

//! [`Server`] — the builder a user of this crate actually touches.
//!
//! # What it is
//!
//! A server is a [`RegistryStore`], a `Vec` of tool handlers, and the
//! negotiation inputs a connection will need. It is a value, not a running
//! process: nothing here binds a port or spawns a task. T13 and T14 supply the
//! transports that take one of these and give it somewhere to run, and until
//! then a test can drive it in-process.
//!
//! # Why a builder and not a `new`
//!
//! A server needs a name before it needs anything else, and most of what
//! follows is optional. A `Server::new(name, tools, resources, prompts)` is
//! four positional arguments in an order nobody can hold in their head, and it
//! is wrong in five of the six shapes a caller wants. The builder makes the
//! order free and names each thing at the call site.

use std::num::NonZeroUsize;
use std::sync::Arc;

use crate::adapters::secondary::idempotency::InMemoryIdempotencyStore;
use crate::application::idempotency::IdempotencyStore;
use crate::application::mcp::input_required::StateKey;
use crate::application::registry::store::RegistryStore;
use crate::application::server::context::RequestContext;
use crate::application::server::live::{Handlers, Live, ServerHandle, ToolEntry};
use crate::domain::ir::JsonObject;
use crate::domain::registry::{
    Primitive, PrimitiveId, PrimitiveKind, Prompt, Registered, Resource, ResourceTemplate, Tool,
};

// `DispatchError` lives in `error` now — it is one of the two error types the
// server API returns, and keeping the pair together is the point of that
// module. Re-exported here because it was defined here: `Server::call` returns
// it and every caller reaching for the return type names it through this
// module, so `server::builder::DispatchError` has to keep resolving.
pub use super::error::DispatchError;

/// A tool handler, and what it answers with.
///
/// The boxed form rather than a generic parameter, deliberately: a server holds
/// a heterogeneous set of these behind one name, so there is nothing to
/// monomorphise over. The alternative — a generic `Server<H>` — buys nothing
/// here and makes the struct unnameable in a field.
///
/// The [`Send`] bound is load-bearing. A handler runs on the runtime, and a
/// handler that is not `Send` cannot be sent to it.
pub type ToolHandler = Arc<dyn Fn(RequestContext, JsonObject) -> BoxFuture + Send + Sync>;

/// The future a [`ToolHandler`] returns.
///
/// Its own alias because spelling `Pin<Box<dyn Future<...> + Send>>` at every
/// call site is how one of them ends up without the `Send`.
pub type BoxFuture = std::pin::Pin<Box<dyn Future<Output = JsonObject> + Send>>;

/// The future a [`ResourceHandler`] or [`PromptHandler`] returns.
///
/// Fallsible where [`BoxFuture`] is not, because a read is allowed to fail for
/// a reason the peer should see as a refusal: a resource whose backing store
/// is unreachable is not an empty resource.
pub type BoxReadFuture =
    std::pin::Pin<Box<dyn Future<Output = Result<JsonObject, DispatchError>> + Send>>;

/// Serves one `resources/read`, given the URI the peer asked for.
///
/// Takes the URI rather than closing over it, because the URI is the request.
/// A handler that closed over its own registration would have to be one per
/// resource, and the point of the lookup table is that the dispatcher finds the
/// right one from the request the peer actually sent.
///
/// Takes the URI by value rather than by reference. A reference would put a
/// lifetime on the handler type, and a closure whose signature carries one
/// cannot be inferred as higher-ranked — so the coercion to this type fails
/// with a lifetime error aimed at the `#[server]` attribute. The dispatcher
/// already owns the string it parsed out of the request, so this costs one
/// clone and buys a handler type with no lifetimes in it at all.
pub type ResourceHandler = Arc<dyn Fn(RequestContext, String) -> BoxReadFuture + Send + Sync>;

/// Serves one `prompts/get`, given the `arguments` the peer sent.
///
/// Takes the whole argument object rather than a decoded one: prompt arguments
/// are free-form by revision — a template like `{diff}` or `{focus}` is the
/// specification's own idiom — so there is no fixed struct to decode into
/// without turning a user's template into this crate's invention.
/// Owned arguments, for the same reason as [`ResourceHandler`].
pub type PromptHandler = Arc<dyn Fn(RequestContext, JsonObject) -> BoxReadFuture + Send + Sync>;

/// A server, built.
///
/// # Examples
///
/// ```
/// use std::sync::Arc;
///
/// use protocol_droid::application::server::Server;
/// use protocol_droid::domain::ir::JsonObject;
///
/// let server = Server::builder("demo")
///     .simple_tool(
///         "echo",
///         r#"{"type":"object","properties":{"text":{"type":"string"}}}"#,
///         Arc::new(|_ctx, args| {
///             Box::pin(async move {
///                 // A handler answers with the tool's result object, not with
///                 // a `Result`: MCP reports a tool failure as a successful
///                 // call carrying `isError: true`, so "the tool failed" is
///                 // something the handler *returns* rather than propagates.
///
///                 // The result is built from a `serde_json::Value` so the
///                 // peer's text is escaped by the serialiser. Hand-writing
///                 // the document would make injection a matter of how many
///                 // quotes a caller remembered.
///                 let echoed = args.as_str().to_string();
///                 let result = JsonObject::from_value(serde_json::json!({
///                     "content": [{ "type": "text", "text": echoed }],
///                 }));
///                 result.unwrap_or_default()
///             })
///         }),
///     )
///     .expect("the schema above is valid")
///     .build();
///
/// assert_eq!(server.tool_count(), 1);
/// assert_eq!(server.name(), "demo");
/// assert!(server.handler("echo").is_some());
/// assert!(server.handler("missing").is_none());
/// ```
// Manual `Debug` rather than derived, for the reason `Dispatcher` has one: the
// store is a `dyn` port with no readable state, so the derived form would
// print `idempotency: <unknown>` on every line of a post-mortem and say
// nothing. What a reader wants is the name, the tool count, and whether a
// store was ever supplied — the third because a server still on the default
// has an in-memory bound, and that is the question to ask first when a
// recorded result has gone missing across a restart.
pub struct Server {
    name: String,
    version: String,
    /// The registry, the handler table, and who is told when they change —
    /// shared with every [`ServerHandle`] so a change made while serving lands
    /// in the same place requests read from.
    live: Arc<Live>,
    /// Deduplicates `tools/call` by request.
    ///
    /// A trait object rather than a concrete store because a caller serving
    /// over HTTP wants a bound that outlives one process — a retried call may
    /// arrive at a different replica — and a concrete `InMemoryIdempotencyStore`
    /// in this field would be the kind of type that compiles everywhere and is
    /// only correct on the one host. The default is the in-memory store, which
    /// is correct for a single process and honest about being so.
    idempotency: Arc<dyn IdempotencyStore<Result = JsonObject>>,
    /// The type-keyed state handlers read through
    /// [`RequestContext::state`](crate::application::server::context::RequestContext::state).
    ///
    /// `Arc` so a `Server` is cheap to clone into a `Dispatcher`, and
    /// the dispatcher is what the transport's per-request handler
    /// actually reaches into. An empty map is the default; a server
    /// that never called [`ServerBuilder::with_state`] gets an empty
    /// map and `state::<T>()` returns `None` for every `T`.
    state: Arc<crate::application::server::context::StateMap>,
    /// What SEP-2322 `requestState` is sealed under. See
    /// [`Server::request_state_key`].
    state_key: Option<StateKey>,
}

impl core::fmt::Debug for Server {
    /// Prints the name, the version, and the tool count.
    ///
    /// Not the store, and the omission is deliberate: a built server always
    /// has one, so printing that would say nothing. Whether a store was
    /// *supplied* is a question about the builder, and the builder's own
    /// `Debug` is where it is answerable.
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Server")
            .field("name", &self.name)
            .field("version", &self.version)
            .field("tool_count", &self.tool_count())
            .finish_non_exhaustive()
    }
}

impl Server {
    /// Serves this server on stdin and stdout until the peer closes stdin.
    ///
    /// The shortest complete serving path in the crate, and the one the
    /// documentation elsewhere assumes exists. The revision is settled from the
    /// peer's own `initialize`, so this neither takes a version nor forces the
    /// caller to pick one — which is the whole reason the crate exists.
    ///
    /// The caller does not have to know that a [`Dispatcher`] and a
    /// [`PeerSession`] are involved, or that the session's profile is resolved
    /// before the loop starts. Both stay reachable: this is a façade over the
    /// types above, not a replacement for them.
    ///
    /// [`Dispatcher`]: crate::application::server::dispatch::Dispatcher
    /// [`PeerSession`]: crate::application::server::context::PeerSession
    ///
    /// # Errors
    ///
    /// Propagates a [`ServeError`] for a read, write, or diagnostics failure.
    /// A protocol error is not one of those: a peer whose frame this server
    /// cannot serve is answered on the wire rather than raised here.
    ///
    /// [`ServeError`]: crate::adapters::primary::stdio::ServeError
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use protocol_droid::application::server::Server;
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let server = Server::builder("greeter").build();
    /// let report = server.serve_stdio().await?;
    /// eprintln!("{report:?}");
    /// # Ok(())
    /// # }
    /// ```
    #[cfg(feature = "stdio")]
    pub async fn serve_stdio(
        self,
    ) -> Result<
        crate::adapters::primary::stdio::ServeReport,
        crate::adapters::primary::stdio::ServeError,
    > {
        let mut loop_ = crate::application::server::serve::StdioBuilder::new(self).build();
        loop_.serve().await
    }

    /// Serves this server over Streamable HTTP, on loopback, until it stops.
    ///
    /// The same one-call shape as [`serve_stdio`](Self::serve_stdio), for the
    /// same reason. The alternative — building an `McpRouter`, an
    /// `axum::Router`, and a listener — moves the work onto the consumer
    /// rather than removing it, and the pieces a consumer would have to reach
    /// for are a `SessionStore` with no production `IdSource` in it and a
    /// `RequestHandler` with no implementation for this `Dispatcher` anywhere
    /// but the tests.
    ///
    /// Binds `127.0.0.1:3000` and serves plain HTTP. That is a development
    /// shape, not a deployment one: a non-loopback address and TLS in front of
    /// it are the caller's to arrange, and [`HttpBuilder`] is how.
    ///
    /// # Errors
    ///
    /// [`ServeError`] from building the route table, from binding, or from the
    /// accept loop. A protocol error is not one of those: a peer whose request
    /// this server cannot serve is answered on the wire rather than raised
    /// here.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use protocol_droid::application::server::Server;
    /// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
    /// let server = Server::builder("greeter").build();
    /// let report = server.serve_http().await?;
    /// eprintln!("listening on {}", report.addr);
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// [`HttpBuilder`]: crate::adapters::primary::transport::http::HttpBuilder
    /// [`ServeError`]: crate::adapters::primary::transport::http::ServeError
    #[cfg(feature = "http")]
    pub async fn serve_http(
        self,
    ) -> Result<
        crate::adapters::primary::transport::http::ServeReport,
        crate::adapters::primary::transport::http::ServeError,
    > {
        crate::adapters::primary::transport::http::HttpBuilder::new(self)
            .serve()
            .await
    }

    /// Starts a server with a name and nothing else.
    ///
    /// Every method on the returned [`ServerBuilder`] is optional, so this
    /// alone is a valid — if useless — server. That is the honest minimum: a
    /// server with no tools is a real state, not a mistake to prevent at the
    /// type level.
    #[must_use]
    pub fn builder(name: impl Into<String>) -> ServerBuilder {
        ServerBuilder {
            name: name.into(),
            version: env!("CARGO_PKG_VERSION").to_string(),
            registry: RegistryStore::new(),
            tools: Vec::new(),
            resource_reads: Vec::new(),
            prompt_gets: Vec::new(),
            idempotency: None,
            state: Vec::new(),
            state_key: None,
        }
    }

    /// The name this server was built with.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The version this server reports in `serverInfo`.
    ///
    /// Free-form on the wire — the schema types it as a plain `string` — so
    /// nothing constrains the format, only that a peer can read it. It lives
    /// beside [`name`](Self::name) because the two are the same fact: they are
    /// the two fields of the `serverInfo` object, they are both about who this
    /// server is rather than what it does, and splitting them would mean a
    /// server whose `name` came from the builder and whose `version` came from
    /// somewhere else.
    #[must_use]
    pub fn version(&self) -> &str {
        &self.version
    }

    /// How many tools are registered, hidden or not.
    #[must_use]
    pub fn tool_count(&self) -> usize {
        self.live.tool_count()
    }

    /// The key `requestState` is sealed under, if one was set.
    ///
    /// Set by [`ServerBuilder::request_state_key`], or filled from the OS
    /// random source by a transport when the server is served without one. A
    /// server with no key can still answer every request; it cannot hand out a
    /// `requestState`, and a retry that carries one is refused as tampered,
    /// because there is nothing to verify it against.
    #[must_use]
    pub const fn request_state_key(&self) -> Option<&StateKey> {
        self.state_key.as_ref()
    }

    /// Installs `key` if the server has none, for a transport that is about to
    /// serve it.
    pub fn ensure_request_state_key(&mut self, key: impl FnOnce() -> Option<StateKey>) {
        if self.state_key.is_none() {
            self.state_key = key();
        }
    }

    /// A handle that adds and removes tools, resources and prompts while this
    /// server is serving, announcing each change to the sessions that
    /// negotiated the matching `listChanged` capability.
    #[must_use]
    pub fn handle(&self) -> ServerHandle {
        ServerHandle::new(Arc::clone(&self.live))
    }

    /// The store `tools/call` is deduplicated through.
    ///
    /// Exposed so a transport can answer "have you seen this id before?"
    /// without running the tool, which is what a resuming peer wants: it has
    /// dropped a connection and wants to know whether the call it never saw
    /// answered actually happened.
    #[must_use]
    pub const fn idempotency(&self) -> &Arc<dyn IdempotencyStore<Result = JsonObject>> {
        &self.idempotency
    }

    /// Registers a state value after the server is built.
    ///
    /// The builder API is `ServerBuilder::with_state`; this method
    /// exists for callers that built a server through a path that did
    /// not surface the builder (the `#[server]` macro's
    /// `into_server` is the canonical one). Each call appends to
    /// the existing map, so a server that already had state for one
    /// type can grow state for a second.
    ///
    /// `Send + Sync + 'static` is what every other port in this crate
    /// asks for; a pool of non-`Send` handles cannot live in this
    /// map.
    pub fn with_state<T: Send + Sync + 'static>(&mut self, value: T) {
        Arc::make_mut(&mut self.state).insert(
            std::any::TypeId::of::<T>(),
            Arc::new(value) as Arc<dyn std::any::Any + Send + Sync>,
        );
    }

    /// Adds a tool to a server that was already built.
    ///
    /// [`ServerBuilder::tool`] is the builder API and takes precedence. This
    /// exists for the same reason [`with_state`](Self::with_state) does: a
    /// caller who went through `#[server]`'s `into_server` never saw a
    /// builder, and until this there was nothing they could do about it.
    ///
    /// It is not a convenience for schemas a Rust type can derive. `#[tool]`
    /// builds `inputSchema` through `schemars` from the handler's parameters,
    /// and `schemars` cannot express what MCP asks a tool to be able to say:
    /// `$schema`, `$defs` with an `$anchor`, `allOf`/`anyOf`, or `if`/`then`/
    /// `else`. Those are SEP-1613 and SEP-2106 — a tool whose schema another
    /// toolchain generated and shipped as JSON has exactly one way to reach a
    /// peer through this crate, and before this method it had none through
    /// the front door.
    ///
    /// Registering and publishing together, as on the builder: a tool in
    /// `tools/list` with no handler is one a client can see and not call, and
    /// a handler with no entry is one that never runs. There is no state in
    /// between for a caller to be wrong in.
    ///
    /// It announces nothing: `&mut self` means nothing can be connected yet.
    /// To change a server that is already serving, use [`Self::handle`].
    ///
    /// Infallible, unlike [`Tool::new`]: a [`Tool`] has already had its
    /// `inputSchema` parsed and validated, so there is no second parse here to
    /// fail.
    pub fn add_tool(&mut self, tool: Tool, handler: ToolHandler) {
        let name = tool.name.clone();
        let _published = self
            .live
            .add(Registered::from(Primitive::from(tool)), |table| {
                table.tools.retain(|entry| entry.name != name);
                table.tools.push(ToolEntry { name, handler });
            });
    }

    /// The state map handlers read through
    /// [`RequestContext::state`](crate::application::server::context::RequestContext::state).
    ///
    /// A transport hands this to every [`RequestContext`]
    /// it builds. A handler that never calls `with_state` on the
    /// builder gets an empty map here, and `state::<T>()` returns
    /// `None` for every `T`.
    #[must_use]
    pub fn state(&self) -> &Arc<crate::application::server::context::StateMap> {
        &self.state
    }

    /// The registry, for the transport that serves it.
    ///
    /// Takes `&self` and hands out an owned snapshot later, so a handler
    /// publishing a new tool does not need the server borrowed mutably while
    /// a request is in flight.
    #[must_use]
    pub fn registry(&self) -> &RegistryStore {
        self.live.registry()
    }

    /// The handler for `name`, if there is one.
    ///
    /// Owned, because the table it came from can be replaced the moment this
    /// returns: a handle may remove the tool while the caller still holds it.
    #[must_use]
    pub fn handler(&self, name: &str) -> Option<ToolHandler> {
        self.live.tool(name)
    }

    /// The reader for `uri`, if one is registered.
    ///
    /// `None` for a resource that is listed but has no reader, which is a
    /// state [`ServerBuilder::resource_with`] cannot produce — registration
    /// takes both together — and which therefore means the URI is simply not
    /// this server's.
    #[must_use]
    pub fn reader(&self, uri: &str) -> Option<ResourceHandler> {
        self.live.reader(uri)
    }

    /// The getter for `name`, if one is registered.
    ///
    /// `None` under the same reasoning as [`reader`](Self::reader).
    #[must_use]
    pub fn getter(&self, name: &str) -> Option<PromptHandler> {
        self.live.getter(name)
    }

    /// Finds `name` and runs it.
    ///
    /// The lookup and the call together rather than a `handler()` the caller
    /// has to remember to use: a transport and an in-process test that each
    /// wrote their own lookup would drift, and the drift would show up as a
    /// tool that is in the registry and not callable.
    ///
    /// Visibility is decided by the **registry**, and only by the registry. It
    /// used to be a second, cheaper check over a cached `required` set held by
    /// each tool entry — which covered features and not scopes. So a
    /// scope-restricted tool was absent from `tools/list` and still ran on
    /// `tools/call`, which is the existence oracle
    /// [`UnknownTool`](DispatchError::UnknownTool) exists to prevent.
    ///
    /// The cached set is gone rather than kept as a belt-and-braces second
    /// check. Once the registry gate runs first, that check can only ever
    /// agree with it, so it decides nothing — and a check that decides nothing
    /// is a check that can drift into deciding the wrong thing without anyone
    /// noticing. One predicate for the whole surface is the property worth
    /// having.
    ///
    /// # Errors
    ///
    /// [`DispatchError::UnknownTool`] when no tool of that name is registered
    /// **or** when this session may not see it — because it requires a
    /// feature the profile lacks, or a scope the caller does not hold. The two
    /// are deliberately the same answer: a different one for each would let a
    /// caller read the existence of a tool it is not allowed to know about.
    /// The server still logs which it was.
    pub async fn call(
        &self,
        ctx: &RequestContext,
        name: &str,
        args: JsonObject,
    ) -> Result<JsonObject, DispatchError> {
        let handler = self.handler(name).ok_or(DispatchError::UnknownTool)?;

        // The gate, once, for both halves of the surface. After this the tool
        // is known to exist *and* to be visible to this session, so the
        // handler is reached with no further checking to disagree with.
        if self
            .registry()
            .get_visible(
                ctx.peer().auth(),
                &ctx.peer().profile(),
                PrimitiveKind::Tool,
                name,
            )
            .is_none()
        {
            return Err(DispatchError::UnknownTool);
        }

        Ok(handler(ctx.clone(), args).await)
    }
}

/// The builder behind [`Server::builder`].
pub struct ServerBuilder {
    name: String,
    version: String,
    registry: RegistryStore,
    tools: Vec<ToolEntry>,
    /// Resources with their readers, in registration order.
    ///
    /// A `Vec` rather than a map because a server registers a handful of
    /// resources and a linear scan over a few entries beats a hash of a
    /// `PrimitiveId` that has to be computed first.
    resource_reads: Vec<(PrimitiveId, ResourceHandler)>,
    /// Prompts with their getters, in registration order.
    prompt_gets: Vec<(PrimitiveId, PromptHandler)>,
    idempotency: Option<Arc<dyn IdempotencyStore<Result = JsonObject>>>,
    /// The state entries built up through `with_state` calls. Held as a
    /// `Vec<(TypeId, Arc<dyn Any>)>` and re-keyed at build time because
    /// the `TypeId` of `T` is only known to `with_state`, and pushing
    /// into a `HashMap` keyed by a value the caller cannot see is the
    /// wrong shape for a builder.
    state: Vec<(std::any::TypeId, Arc<dyn std::any::Any + Send + Sync>)>,
    state_key: Option<StateKey>,
}

impl core::fmt::Debug for ServerBuilder {
    /// Prints the same three facts as [`Server`], plus whether a store was
    /// supplied.
    ///
    /// The last is the one worth having here, and it is the question to ask
    /// first when a recorded result has gone missing across a restart: a
    /// server still on the default has an in-process bound, and no way to see
    /// that from the built server it became.
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ServerBuilder")
            .field("name", &self.name)
            .field("version", &self.version)
            .field("tool_count", &self.tools.len())
            .field("idempotency_supplied", &self.idempotency.is_some())
            .finish_non_exhaustive()
    }
}

impl ServerBuilder {
    /// Sets the key SEP-2322 `requestState` is sealed under.
    ///
    /// Optional: a transport fills one from the OS random source when it is
    /// absent. Set it explicitly when several replicas serve one endpoint —
    /// they must share it, or a retry that lands on a different replica finds
    /// its state refused.
    #[must_use]
    pub const fn request_state_key(mut self, key: StateKey) -> Self {
        self.state_key = Some(key);
        self
    }

    /// Sets the version this server reports in `serverInfo`.
    ///
    /// Defaults to the version of the crate that built it, which is right for
    /// a server that *is* this crate — the examples, the conformance runner —
    /// and wrong for one that merely links it. A library user shipping their
    /// own server should say so here rather than let the dependency's version
    /// stand in for theirs, which would have every downstream server claiming
    /// to be whatever `protocol-droid` happened to be pinned at.
    #[must_use]
    pub fn version(mut self, version: impl Into<String>) -> Self {
        self.version = version.into();
        self
    }

    /// Replaces the store `tools/call` is deduplicated through.
    ///
    /// For a store that outlives the process — a shared cache, a table — or
    /// for a test that wants to assert on what was recorded. A server that
    /// keeps its own `IdSource` across restarts has a gap the in-memory
    /// default cannot close: a call recorded before a restart is unknown
    /// after it, and the retry runs the tool again.
    #[must_use]
    pub fn idempotency(mut self, store: Arc<dyn IdempotencyStore<Result = JsonObject>>) -> Self {
        self.idempotency = Some(store);
        self
    }

    /// Registers a value handlers can read through
    /// [`RequestContext::state`](crate::application::server::context::RequestContext::state).
    ///
    /// `T` is keyed by `TypeId`, so two `with_state` calls with the same
    /// `T` keep the later value (the registry is a `HashMap`, and the
    /// later write is the one the lookup will see). The `Send + Sync`
    /// bounds are the same every other port in this crate asks for; a
    /// pool of non-`Send` handles cannot live in this map.
    ///
    /// The point of this method is the `#[tool]` problem: a macro that
    /// generates module-scope items cannot see `self`, so the handler
    /// it generates has no captured environment. Storing the state
    /// here and reading it through the request context is the
    /// standard answer, and the only one that keeps the macro
    /// out of the receiver business.
    #[must_use]
    pub fn with_state<T: Send + Sync + 'static>(mut self, value: T) -> Self {
        self.state.push((
            std::any::TypeId::of::<T>(),
            Arc::new(value) as Arc<dyn std::any::Any + Send + Sync>,
        ));
        self
    }

    /// Adds a tool, and publishes it to the registry in the same step.
    ///
    /// Takes a whole [`Tool`] rather than a name and a schema string, because
    /// a schema string can only ever express the plainest tool there is. An
    /// output schema, annotations and icons all change what the tool
    /// _requires_ of a session, and a builder that cannot name them would
    /// register a tool whose requirement nothing had computed.
    ///
    /// Registering and publishing together is the point: a tool in the list
    /// but not in the registry is a tool a client can see and not call, and
    /// one in neither is a handler that never runs. There is no state in
    /// between for a caller to be wrong in.
    #[must_use]
    pub fn tool(self, tool: Tool, handler: ToolHandler) -> Self {
        self.restricted_tool(Registered::from(Primitive::from(tool)), handler)
    }

    /// Adds a tool that carries its own requirements.
    ///
    /// [`tool`](Self::tool) takes a [`Tool`], and a `Tool` cannot say
    /// "callers need the `files:read` scope" — that lives on [`Registered`],
    /// which wraps a primitive with the policy it must satisfy. So without this
    /// method the scope half of `is_visible` is reachable from no builder
    /// call, `auth.rs` is dead code, and a server can be built that has no way
    /// to express an authorisation rule at all. That is the same dead-gate
    /// shape T11 found, one layer up.
    ///
    /// A `Registered` carrying something other than a tool is registered
    /// under its own id and left out of the call table: it stays visible to
    /// `tools/list` and the listing is the truth, so nothing is silently
    /// dropped, but there is no handler to name for a `tools/call` of it.
    #[must_use]
    pub fn restricted_tool(mut self, registered: Registered, handler: ToolHandler) -> Self {
        if let Primitive::Tool(tool) = registered.primitive() {
            self.tools.push(ToolEntry {
                name: tool.name.clone(),
                handler,
            });
        }
        self.registry.register(registered);
        self
    }

    /// Adds a plain tool: a name, an input schema, and a handler.
    ///
    /// The common case, and the one a tool without an output schema or
    /// annotations fits. It exists so the common case is one line rather than
    /// a `Tool::new(...)?` at every call site, and it returns a `Result`
    /// because the schema is the one fallible step.
    /// # Errors
    ///
    /// The [`JsonObject`](crate::domain::ir::JsonObjectError) from parsing
    /// `input_schema`. The schema is validated here, once, so a malformed one
    /// fails at build time rather than on the first call.
    pub fn simple_tool(
        self,
        name: impl Into<PrimitiveId>,
        input_schema: &str,
        handler: ToolHandler,
    ) -> Result<Self, crate::domain::ir::JsonObjectError> {
        Ok(self.tool(Tool::new(name, input_schema)?, handler))
    }

    /// Adds a resource.
    ///
    /// Adds a resource together with the reader that answers `resources/read`
    /// for it.
    ///
    /// Both at once, and there is deliberately no `resource(resource)` beside
    /// it. A resource that is listed by `resources/list` and then refused by
    /// `resources/read` is a server telling a peer something false, and the
    /// dispatcher has to answer `InvalidParams` for the read because it has no
    /// other way to distinguish that from a URI the server never had. Making
    /// the reader a separate opt-in would put that state one method call away,
    /// and the failure would surface to a client rather than at the call site
    /// that built it.
    #[must_use]
    pub fn resource_with(self, resource: Resource, reader: ResourceHandler) -> Self {
        let mut builder = self;
        builder.resource_reads.push((resource.uri.clone(), reader));
        builder
            .registry
            .register(Registered::from(Primitive::from(resource)));
        builder
    }

    /// Adds a resource template.
    ///
    /// Templates are listed by their own method, and they are a distinct
    /// [`PrimitiveKind`] in the registry, so a template needs its own entry
    /// point rather than sharing `resource`'s.
    #[must_use]
    pub fn resource_template(self, template: ResourceTemplate) -> Self {
        self.registry
            .register(Registered::from(Primitive::from(template)));
        self
    }

    /// Adds a prompt together with the getter that answers `prompts/get` for it.
    ///
    /// Required together for the reason [`resource_with`](Self::resource_with)
    /// takes both.
    #[must_use]
    pub fn prompt_with(self, prompt: Prompt, getter: PromptHandler) -> Self {
        let mut builder = self;
        builder.prompt_gets.push((prompt.name.clone(), getter));
        builder
            .registry
            .register(Registered::from(Primitive::from(prompt)));
        builder
    }

    /// Finishes the server.
    #[must_use]
    pub fn build(self) -> Server {
        // Deduplicating by default rather than on request. A caller who has
        // never heard of idempotency is the caller most likely to be serving
        // a retriable HTTP transport, where a dropped connection turns one
        // call into two, and a store that has to be opted into is a store
        // nobody will think to opt into. The default is the bound that fails
        // safe: a server that forgets runs a tool once rather than twice.
        let idempotency = self.idempotency.unwrap_or_else(|| {
            // 4096 keys is generous for a single process and small enough to
            // bound memory. The capacity is a policy number, not a constant
            // with a meaning, so it is named rather than inlined.
            const DEFAULT_KEYS: NonZeroUsize = match NonZeroUsize::new(4096) {
                Some(value) => value,
                // Unreachable: 4096 is not zero. A `match` rather than an
                // `unwrap` because the crate denies that.
                None => NonZeroUsize::MIN,
            };
            Arc::new(InMemoryIdempotencyStore::with_capacity(DEFAULT_KEYS))
        });
        Server {
            name: self.name,
            version: self.version,
            live: Arc::new(Live::new(
                self.registry,
                Handlers {
                    tools: self.tools,
                    resources: self.resource_reads,
                    prompts: self.prompt_gets,
                },
            )),
            idempotency,
            state: Arc::new(
                self.state
                    .into_iter()
                    .collect::<std::collections::HashMap<_, _>>(),
            ),
            state_key: self.state_key,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::{DispatchError, Server};
    use crate::application::server::context::{PeerSession, RequestContext};
    use crate::application::session::SinkError;
    use crate::domain::feature::FeatureProfile;
    use crate::domain::ir::{JsonObject, JsonObjectError};
    use crate::domain::registry::PrimitiveKind;
    use crate::domain::session::{CancellationToken, SessionId};
    use crate::domain::{ProtocolVersion, RequestId, TransportKind};

    /// What a test here can fail with.
    #[derive(Debug)]
    enum Failure {
        /// The test's own complaint.
        Message(String),
        /// A fixture was not a JSON object.
        Fixture(JsonObjectError),
    }

    impl core::fmt::Display for Failure {
        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
            match self {
                Self::Message(text) => f.write_str(text),
                Self::Fixture(error) => write!(f, "bad fixture: {error}"),
            }
        }
    }

    impl std::error::Error for Failure {}

    impl From<String> for Failure {
        fn from(text: String) -> Self {
            Self::Message(text)
        }
    }

    impl From<&str> for Failure {
        fn from(text: &str) -> Self {
            Self::Message(text.to_string())
        }
    }

    impl From<JsonObjectError> for Failure {
        fn from(error: JsonObjectError) -> Self {
            Self::Fixture(error)
        }
    }

    type TestError = Failure;

    /// A handler that answers with a fixed object, so a test can tell one tool
    /// from another without every case writing its own closure.
    fn answering(answer: &str) -> Result<super::ToolHandler, TestError> {
        let answer = JsonObject::new(answer)?;
        Ok(Arc::new(move |_ctx: RequestContext, _args: JsonObject| {
            let answer = answer.clone();
            Box::pin(async move { answer })
        }))
    }

    /// A port pair that answers nothing and records nothing.
    ///
    /// The builder tests never elicit, sample, or emit progress — they are
    /// about registration and dispatch — so the ports exist only to satisfy
    /// `PeerSession`. A `RefCell` is safe here precisely because nothing
    /// touches them, and a real channel would be a second thing to get wrong
    /// in a test that is not about channels.
    struct Inert;

    impl crate::application::session::FrameSink for Inert {
        fn send(&self, _frame: &str) -> impl Future<Output = Result<(), SinkError>> + Send {
            std::future::ready(Ok(()))
        }
    }

    impl crate::application::server::context::CallbackExchange for Inert {
        fn exchange(
            &self,
            _method: &str,
            _params: &JsonObject,
        ) -> impl Future<Output = Result<JsonObject, crate::application::server::ExchangeError>> + Send
        {
            std::future::ready(Err(crate::application::server::ExchangeError::new(
                "inert in these tests",
            )))
        }
    }

    fn peer_at(profile: FeatureProfile) -> PeerSession {
        PeerSession::new(
            ProtocolVersion::V2025_06_18,
            profile,
            SessionId::new(1),
            TransportKind::Stdio,
            Arc::new(Inert),
            Arc::new(Inert),
        )
    }

    fn ctx_for(profile: FeatureProfile) -> RequestContext {
        RequestContext::for_request(
            &peer_at(profile),
            RequestId::number(1),
            CancellationToken::new(),
            None,
        )
    }

    /// A tool added to the builder is in the registry in the same step.
    ///
    /// The builder's claim is that there is no state in between — no tool that
    /// is listable but not callable. A registry the builder has to be told
    /// about separately cannot make that claim, so this is the test of the
    /// claim and not a restatement of the code.
    #[test]
    fn a_registered_tool_is_in_the_registry_immediately() -> Result<(), TestError> {
        let server = Server::builder("demo")
            .simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
            .build();

        assert_eq!(server.tool_count(), 1);
        assert_eq!(server.registry().snapshot().len(PrimitiveKind::Tool), 1);
        Ok(())
    }

    /// A tool added after `into_server` lists *and* runs.
    ///
    /// Both halves, because the registry and the handler table are two places
    /// a method like this could write to only one of. A tool in `tools/list`
    /// with no handler is one a client can see and not call — the failure mode
    /// the builder's own doc calls out — and it would pass a listing-only test.
    #[test]
    fn a_tool_added_after_the_build_lists_and_runs() -> Result<(), TestError> {
        use crate::domain::registry::{Primitive, Tool};

        let mut server = Server::builder("demo")
            .simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
            .build();

        // The schema is the one thing `add_tool` exists for: `$defs`, an
        // `$anchor` and `if`/`then`/`else` are JSON Schema 2020-12 keywords no
        // Rust type derives, so this is the whole reason the method is here.
        let raw = Tool::new(
            "json_schema_2020_12",
            r#"{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object","$defs":{"address":{"$anchor":"addressDef","type":"string"}},"if":{"properties":{"a":{"const":1}}},"then":{"required":["b"]}}"#,
        )?;
        server.add_tool(raw, answering(r#"{"ok":true}"#)?);

        assert_eq!(server.tool_count(), 2);
        assert_eq!(server.registry().snapshot().len(PrimitiveKind::Tool), 2);
        assert!(server.handler("json_schema_2020_12").is_some());

        // The listing, not just the count: the schema has to reach `tools/list`
        // with the keywords intact, since that is the wire the suite reads.
        let snapshot = server.registry().snapshot();
        let advertised = snapshot
            .list(PrimitiveKind::Tool)
            .iter()
            .find(|tool| tool.id().as_str() == "json_schema_2020_12")
            .ok_or("the added tool is not in the registry")?;
        let Primitive::Tool(advertised) = advertised.primitive() else {
            return Err("a tool listing held something that is not a tool".into());
        };
        let text = advertised.input_schema.as_str();
        for keyword in ["$schema", "$defs", "$anchor", "if", "then"] {
            if !text.contains(keyword) {
                return Err(format!("{keyword} did not reach the listing").into());
            }
        }
        Ok(())
    }

    /// A malformed input schema fails at build time, not on the first call.
    #[test]
    fn a_malformed_input_schema_is_rejected_at_build_time() -> Result<(), TestError> {
        // The assertion is on the `Result` itself rather than on an error
        // message: the error's text is not the contract, the timing of the
        // failure is. A test that unwrapped and matched on a string would pass
        // against a builder that validated for the wrong reason.
        let handler = answering("{}")?;
        let built = Server::builder("demo").simple_tool("bad", "{not json", handler);

        match built {
            Ok(_) => Err(Failure::from(
                "a malformed input schema was accepted at build time",
            )),
            Err(_) => Ok(()),
        }
    }

    /// Two tools with one name is one tool, not two.
    ///
    /// T11's store treats a re-registration as a replacement, and the builder
    /// has to agree or the registry would list a name the dispatcher cannot
    /// resolve to a single handler.
    #[test]
    fn a_second_tool_of_the_same_name_replaces_the_first() -> Result<(), TestError> {
        let server = Server::builder("demo")
            .simple_tool(
                "echo",
                r#"{"type":"object"}"#,
                answering(r#"{"first":true}"#)?,
            )?
            .simple_tool(
                "echo",
                r#"{"type":"object"}"#,
                answering(r#"{"second":true}"#)?,
            )?
            .build();

        assert_eq!(server.tool_count(), 2, "both handlers are kept");
        assert_eq!(
            server.registry().snapshot().len(PrimitiveKind::Tool),
            1,
            "but the registry holds one name"
        );
        Ok(())
    }

    /// An unknown tool is refused without reaching any handler.
    #[tokio::test]
    async fn an_unknown_tool_is_refused() -> Result<(), TestError> {
        let server = Server::builder("demo")
            .simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
            .build();
        let ctx = ctx_for(FeatureProfile::all());

        match server.call(&ctx, "nope", JsonObject::new("{}")?).await {
            Ok(answer) => Err(Failure::from(format!("expected a refusal, got {answer}"))),
            Err(DispatchError::UnknownTool) => Ok(()),
            Err(other) => Err(Failure::from(format!(
                "refused for the wrong reason: {other}"
            ))),
        }
    }

    /// A registered tool whose required features the session lacks is refused
    /// by `call`, not run and refused afterwards.
    ///
    /// A tool with an output schema requires `ToolStructuredOutput`. A session
    /// that never negotiated it must not run the tool and then drop the part
    /// of the answer the client could not have used.
    #[tokio::test]
    async fn a_tool_needing_an_absent_feature_is_not_run() -> Result<(), TestError> {
        use crate::domain::registry::Tool;

        // The full `Tool`, not `simple_tool`: an output schema is what raises
        // the requirement in the first place, and `simple_tool` could not
        // express one. A test written against it would be testing a gap.
        let structured = Tool::new("echo", r#"{"type":"object"}"#)?
            .with_output_schema(r#"{"type":"object"}"#)?;

        let server = Server::builder("demo")
            .tool(structured, answering(r#"{"ran":true}"#)?)
            .build();

        let ctx = ctx_for(FeatureProfile::empty());
        match server.call(&ctx, "echo", JsonObject::new("{}")?).await {
            Ok(answer) => Err(Failure::from(format!(
                "a tool the session cannot call returned {answer}"
            ))),
            // `UnknownTool`, not `NotCallable`, and the distinction is the
            // point. A tool this session may not see is a tool this session is
            // not told about: answering "it exists but you may not have it"
            // would confirm it exists, and a caller could then enumerate the
            // registry by watching which names come back refused. The same
            // answer as a tool that is not there is what makes the refusal
            // probe-proof.
            Err(DispatchError::UnknownTool) => Ok(()),
            Err(other) => Err(Failure::from(format!(
                "refused for the wrong reason: {other}"
            ))),
        }
    }
}