Skip to main content

tower_mcp/
tool.rs

1//! Tool definition and builder API
2//!
3//! Provides ergonomic ways to define MCP tools:
4//!
5//! 1. **Builder pattern** - Fluent API for defining tools
6//! 2. **Trait-based** - Implement `McpTool` for full control
7//! 3. **Function-based** - Quick tools from async functions
8//!
9//! ## Per-Tool Middleware
10//!
11//! Tools are implemented as Tower services internally, enabling middleware
12//! composition via the `.layer()` method:
13//!
14//! ```rust
15//! use std::time::Duration;
16//! use tower::timeout::TimeoutLayer;
17//! use tower_mcp::{ToolBuilder, CallToolResult};
18//! use schemars::JsonSchema;
19//! use serde::Deserialize;
20//!
21//! #[derive(Debug, Deserialize, JsonSchema)]
22//! struct SearchInput { query: String }
23//!
24//! let tool = ToolBuilder::new("slow_search")
25//!     .description("Search with extended timeout")
26//!     .handler(|input: SearchInput| async move {
27//!         Ok(CallToolResult::text("result"))
28//!     })
29//!     .layer(TimeoutLayer::new(Duration::from_secs(30)))
30//!     .build();
31//! ```
32
33use std::borrow::Cow;
34use std::convert::Infallible;
35use std::fmt;
36use std::future::Future;
37use std::pin::Pin;
38use std::sync::Arc;
39use std::task::{Context, Poll};
40
41use pin_project_lite::pin_project;
42
43use schemars::{JsonSchema, Schema, SchemaGenerator};
44use serde::Serialize;
45use serde::de::DeserializeOwned;
46use serde_json::{Map, Value};
47#[cfg(feature = "stateless")]
48use tower::ServiceExt;
49use tower::util::BoxCloneService;
50use tower_service::Service;
51
52#[cfg(feature = "stateless")]
53use tokio::sync::Mutex;
54
55use crate::context::{Extensions, RequestContext};
56use crate::error::{Error, Result, ResultExt};
57use crate::protocol::{
58    CallToolResult, ClientCapabilities, RequestOutcome, TaskSupportMode, ToolAnnotations,
59    ToolDefinition, ToolExecution, ToolIcon,
60};
61
62// =============================================================================
63// Service Types for Per-Tool Middleware
64// =============================================================================
65
66/// Request type for tool services.
67///
68/// Contains the request context (for progress reporting, cancellation, etc.)
69/// and the tool arguments as raw JSON.
70#[derive(Debug, Clone)]
71pub struct ToolRequest {
72    /// Request context for progress reporting, cancellation, and client requests
73    pub ctx: RequestContext,
74    /// Tool arguments as raw JSON
75    pub args: Value,
76}
77
78impl ToolRequest {
79    /// Create a new tool request
80    pub fn new(ctx: RequestContext, args: Value) -> Self {
81        Self { ctx, args }
82    }
83}
84
85/// A boxed, cloneable tool service with `Error = Infallible`.
86///
87/// This is the internal service type that tools use. Middleware errors are
88/// caught and converted to `CallToolResult::error()` responses, so the
89/// service never fails at the Tower level.
90pub type BoxToolService = BoxCloneService<ToolRequest, CallToolResult, Infallible>;
91
92/// A boxed MRTR-capable tool service.
93#[cfg(feature = "stateless")]
94type BoxMrtrToolService = BoxCloneService<ToolRequest, RequestOutcome<CallToolResult>, Infallible>;
95
96/// Catches errors from the inner service and converts them to `CallToolResult::error()`.
97///
98/// This wrapper ensures that middleware errors (e.g., timeouts, rate limits)
99/// and handler errors are converted to tool-level error responses with
100/// `is_error: true`, rather than propagating as Tower service errors.
101#[doc(hidden)]
102pub struct ToolCatchError<S> {
103    inner: S,
104}
105
106impl<S> ToolCatchError<S> {
107    /// Create a new `ToolCatchError` wrapping the given service.
108    pub fn new(inner: S) -> Self {
109        Self { inner }
110    }
111}
112
113impl<S: Clone> Clone for ToolCatchError<S> {
114    fn clone(&self) -> Self {
115        Self {
116            inner: self.inner.clone(),
117        }
118    }
119}
120
121impl<S: fmt::Debug> fmt::Debug for ToolCatchError<S> {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        f.debug_struct("ToolCatchError")
124            .field("inner", &self.inner)
125            .finish()
126    }
127}
128
129pin_project! {
130    /// Future for [`ToolCatchError`].
131    #[doc(hidden)]
132    pub struct ToolCatchErrorFuture<F> {
133        #[pin]
134        inner: F,
135    }
136}
137
138impl<F, E> Future for ToolCatchErrorFuture<F>
139where
140    F: Future<Output = std::result::Result<CallToolResult, E>>,
141    E: fmt::Display,
142{
143    type Output = std::result::Result<CallToolResult, Infallible>;
144
145    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
146        match self.project().inner.poll(cx) {
147            Poll::Pending => Poll::Pending,
148            Poll::Ready(Ok(result)) => Poll::Ready(Ok(result)),
149            Poll::Ready(Err(err)) => Poll::Ready(Ok(CallToolResult::error(err.to_string()))),
150        }
151    }
152}
153
154impl<S> Service<ToolRequest> for ToolCatchError<S>
155where
156    S: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
157    S::Error: fmt::Display + Send,
158    S::Future: Send,
159{
160    type Response = CallToolResult;
161    type Error = Infallible;
162    type Future = ToolCatchErrorFuture<S::Future>;
163
164    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
165        // Map any readiness error to Infallible (we catch it on call)
166        match self.inner.poll_ready(cx) {
167            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
168            Poll::Ready(Err(_)) => Poll::Ready(Ok(())),
169            Poll::Pending => Poll::Pending,
170        }
171    }
172
173    fn call(&mut self, req: ToolRequest) -> Self::Future {
174        ToolCatchErrorFuture {
175            inner: self.inner.call(req),
176        }
177    }
178}
179
180/// Catches errors from an MRTR-capable tool service.
181///
182/// Per-tool middleware has the same error semantics for complete and MRTR
183/// handlers: middleware and handler failures become complete tool error
184/// results, while input-required outcomes pass through unchanged.
185#[cfg(feature = "stateless")]
186#[derive(Clone)]
187struct MrtrToolCatchError<S> {
188    inner: S,
189}
190
191#[cfg(feature = "stateless")]
192impl<S> MrtrToolCatchError<S> {
193    fn new(inner: S) -> Self {
194        Self { inner }
195    }
196}
197
198#[cfg(feature = "stateless")]
199impl<S> Service<ToolRequest> for MrtrToolCatchError<S>
200where
201    S: Service<ToolRequest, Response = RequestOutcome<CallToolResult>> + Clone + Send + 'static,
202    S::Error: fmt::Display + Send + 'static,
203    S::Future: Send + 'static,
204{
205    type Response = RequestOutcome<CallToolResult>;
206    type Error = Infallible;
207    type Future =
208        Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
209
210    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
211        match self.inner.poll_ready(cx) {
212            Poll::Ready(Ok(())) | Poll::Ready(Err(_)) => Poll::Ready(Ok(())),
213            Poll::Pending => Poll::Pending,
214        }
215    }
216
217    fn call(&mut self, req: ToolRequest) -> Self::Future {
218        let future = self.inner.call(req);
219        Box::pin(async move {
220            Ok(match future.await {
221                Ok(outcome) => outcome,
222                Err(error) => RequestOutcome::Complete(CallToolResult::error(error.to_string())),
223            })
224        })
225    }
226}
227
228/// A tower [`Layer`](tower::Layer) that applies a guard function before the inner service.
229///
230/// Guards run before the tool handler and can short-circuit with an error message.
231/// Use via [`ToolBuilderWithHandler::guard`] or [`Tool::with_guard`] rather than
232/// constructing directly.
233///
234/// # Example
235///
236/// ```rust
237/// use tower_mcp::{ToolBuilder, ToolRequest, CallToolResult};
238/// use schemars::JsonSchema;
239/// use serde::Deserialize;
240///
241/// #[derive(Debug, Deserialize, JsonSchema)]
242/// struct DeleteInput { id: String, confirm: bool }
243///
244/// let tool = ToolBuilder::new("delete")
245///     .description("Delete a record")
246///     .handler(|input: DeleteInput| async move {
247///         Ok(CallToolResult::text(format!("deleted {}", input.id)))
248///     })
249///     .guard(|req: &ToolRequest| {
250///         let confirm = req.args.get("confirm").and_then(|v| v.as_bool()).unwrap_or(false);
251///         if !confirm {
252///             return Err("Must set confirm=true to delete".to_string());
253///         }
254///         Ok(())
255///     })
256///     .build();
257/// ```
258#[derive(Clone)]
259pub struct GuardLayer<G> {
260    guard: G,
261}
262
263impl<G> GuardLayer<G> {
264    /// Create a new guard layer from a closure.
265    ///
266    /// The closure receives a `&ToolRequest` and returns `Ok(())` to proceed
267    /// or `Err(String)` to reject with an error message.
268    pub fn new(guard: G) -> Self {
269        Self { guard }
270    }
271}
272
273impl<G, S> tower::Layer<S> for GuardLayer<G>
274where
275    G: Clone,
276{
277    type Service = GuardService<G, S>;
278
279    fn layer(&self, inner: S) -> Self::Service {
280        GuardService {
281            guard: self.guard.clone(),
282            inner,
283        }
284    }
285}
286
287/// Service wrapper that runs a guard check before calling the inner service.
288///
289/// Created by [`GuardLayer`]. See its documentation for usage.
290#[doc(hidden)]
291#[derive(Clone)]
292pub struct GuardService<G, S> {
293    guard: G,
294    inner: S,
295}
296
297impl<G, S, R> Service<ToolRequest> for GuardService<G, S>
298where
299    G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
300    S: Service<ToolRequest, Response = R> + Clone + Send + 'static,
301    S::Error: Into<Error> + Send,
302    S::Future: Send,
303    R: Send + 'static,
304{
305    type Response = R;
306    type Error = Error;
307    type Future = Pin<Box<dyn Future<Output = std::result::Result<R, Error>> + Send>>;
308
309    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
310        self.inner.poll_ready(cx).map_err(Into::into)
311    }
312
313    fn call(&mut self, req: ToolRequest) -> Self::Future {
314        match (self.guard)(&req) {
315            Ok(()) => {
316                let fut = self.inner.call(req);
317                Box::pin(async move { fut.await.map_err(Into::into) })
318            }
319            Err(msg) => Box::pin(async move { Err(Error::tool(msg)) }),
320        }
321    }
322}
323
324/// A marker type for tools that take no parameters.
325///
326/// Use this instead of `()` when defining tools with no input parameters.
327/// The unit type `()` generates `"type": "null"` in JSON Schema, which many
328/// MCP clients reject. `NoParams` generates `"type": "object"` with no
329/// required properties, which is the correct schema for parameterless tools.
330///
331/// # Example
332///
333/// ```rust
334/// use tower_mcp::{ToolBuilder, CallToolResult, NoParams};
335///
336/// let tool = ToolBuilder::new("get_status")
337///     .description("Get current status")
338///     .handler(|_input: NoParams| async move {
339///         Ok(CallToolResult::text("OK"))
340///     })
341///     .build();
342/// ```
343#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
344pub struct NoParams;
345
346impl<'de> serde::Deserialize<'de> for NoParams {
347    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
348    where
349        D: serde::Deserializer<'de>,
350    {
351        // Accept null, empty object, or any object (ignoring all fields)
352        struct NoParamsVisitor;
353
354        impl<'de> serde::de::Visitor<'de> for NoParamsVisitor {
355            type Value = NoParams;
356
357            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
358                formatter.write_str("null or an object")
359            }
360
361            fn visit_unit<E>(self) -> std::result::Result<Self::Value, E>
362            where
363                E: serde::de::Error,
364            {
365                Ok(NoParams)
366            }
367
368            fn visit_none<E>(self) -> std::result::Result<Self::Value, E>
369            where
370                E: serde::de::Error,
371            {
372                Ok(NoParams)
373            }
374
375            fn visit_some<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
376            where
377                D: serde::Deserializer<'de>,
378            {
379                serde::Deserialize::deserialize(deserializer)
380            }
381
382            fn visit_map<A>(self, mut map: A) -> std::result::Result<Self::Value, A::Error>
383            where
384                A: serde::de::MapAccess<'de>,
385            {
386                // Drain the map, ignoring all entries
387                while map
388                    .next_entry::<serde::de::IgnoredAny, serde::de::IgnoredAny>()?
389                    .is_some()
390                {}
391                Ok(NoParams)
392            }
393        }
394
395        deserializer.deserialize_any(NoParamsVisitor)
396    }
397}
398
399impl JsonSchema for NoParams {
400    fn schema_name() -> Cow<'static, str> {
401        Cow::Borrowed("NoParams")
402    }
403
404    fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
405        serde_json::json!({
406            "type": "object"
407        })
408        .try_into()
409        .expect("valid schema")
410    }
411}
412
413/// Validate a tool name according to MCP spec (SEP-986).
414///
415/// Tool names must be:
416/// - 1-64 characters long
417/// - Contain only ASCII alphanumeric characters, underscores, hyphens, dots,
418///   and forward slashes
419///
420/// Returns `Ok(())` if valid, `Err` with description if invalid.
421pub(crate) fn validate_tool_name(name: &str) -> Result<()> {
422    if name.is_empty() {
423        return Err(Error::tool("Tool name cannot be empty"));
424    }
425    if name.len() > 64 {
426        return Err(Error::tool(format!(
427            "Tool name '{}' exceeds maximum length of 64 characters (got {})",
428            name,
429            name.len()
430        )));
431    }
432    if let Some(invalid_char) = name
433        .chars()
434        .find(|c| !c.is_ascii_alphanumeric() && *c != '_' && *c != '-' && *c != '.' && *c != '/')
435    {
436        return Err(Error::tool(format!(
437            "Tool name '{}' contains invalid character '{}'. Only alphanumeric, underscore, hyphen, dot, and forward slash are allowed.",
438            name, invalid_char
439        )));
440    }
441    Ok(())
442}
443
444/// Ensure a JSON Schema value has `"type": "object"`.
445///
446/// The MCP spec requires tool input schemas to be JSON objects with a `"type"` field.
447/// Some types (e.g., `serde_json::Value`) generate schemas via schemars that lack
448/// the `"type"` field, which causes MCP clients to reject the tool.
449pub(crate) fn ensure_object_schema(mut schema: Value) -> Value {
450    if let Some(obj) = schema.as_object_mut()
451        && !obj.contains_key("type")
452    {
453        obj.insert("type".to_string(), serde_json::json!("object"));
454    }
455    schema
456}
457
458/// A boxed future for tool handlers
459pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
460
461/// Identity allocated for one task-backed tool execution.
462///
463/// The same value is supplied to task preparation and inserted into the
464/// background handler's request extensions.
465#[derive(Debug, Clone, PartialEq, Eq)]
466pub struct TaskContext {
467    task_id: String,
468}
469
470impl TaskContext {
471    pub(crate) fn new(task_id: String) -> Self {
472        Self { task_id }
473    }
474
475    /// The server-generated task identifier.
476    pub fn task_id(&self) -> &str {
477        &self.task_id
478    }
479}
480
481/// Metadata and application state produced before task execution begins.
482#[derive(Debug, Clone, Default)]
483pub struct TaskPreparation {
484    pub(crate) meta: Option<Map<String, Value>>,
485    pub(crate) extensions: Extensions,
486}
487
488impl TaskPreparation {
489    /// Create an empty preparation result.
490    pub fn new() -> Self {
491        Self::default()
492    }
493
494    /// Attach protocol `_meta` to every view of this task.
495    pub fn with_meta(mut self, meta: Map<String, Value>) -> Self {
496        self.meta = Some(meta);
497        self
498    }
499
500    /// Make application state available to the background handler through
501    /// [`crate::extract::Extension`].
502    pub fn with_extension<T: Send + Sync + 'static>(mut self, value: T) -> Self {
503        self.extensions.insert(value);
504        self
505    }
506}
507
508pub(crate) trait TaskPreparer: Send + Sync {
509    fn prepare(
510        &self,
511        context: TaskContext,
512        arguments: Value,
513    ) -> BoxFuture<'_, Result<TaskPreparation>>;
514}
515
516impl<F, Fut> TaskPreparer for F
517where
518    F: Fn(TaskContext, Value) -> Fut + Send + Sync,
519    Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
520{
521    fn prepare(
522        &self,
523        context: TaskContext,
524        arguments: Value,
525    ) -> BoxFuture<'_, Result<TaskPreparation>> {
526        Box::pin((self)(context, arguments))
527    }
528}
529
530struct TypedTaskPreparer<I, F> {
531    prepare: F,
532    _phantom: std::marker::PhantomData<I>,
533}
534
535impl<I, F, Fut> TaskPreparer for TypedTaskPreparer<I, F>
536where
537    I: DeserializeOwned + Send + Sync + 'static,
538    F: Fn(TaskContext, I) -> Fut + Send + Sync,
539    Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
540{
541    fn prepare(
542        &self,
543        context: TaskContext,
544        arguments: Value,
545    ) -> BoxFuture<'_, Result<TaskPreparation>> {
546        let input = serde_json::from_value(arguments)
547            .map_err(|error| Error::invalid_params(format!("Invalid input: {error}")));
548        match input {
549            Ok(input) => Box::pin((self.prepare)(context, input)),
550            Err(error) => Box::pin(async move { Err(error) }),
551        }
552    }
553}
554
555/// Tool handler trait - the core abstraction for tool execution
556pub trait ToolHandler: Send + Sync {
557    /// Execute the tool with the given arguments
558    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>>;
559
560    /// Execute the tool with request context for progress/cancellation support
561    ///
562    /// The default implementation ignores the context and calls `call`.
563    /// Override this to receive progress/cancellation context.
564    fn call_with_context(
565        &self,
566        _ctx: RequestContext,
567        args: Value,
568    ) -> BoxFuture<'_, Result<CallToolResult>> {
569        self.call(args)
570    }
571
572    /// Returns true if this handler uses context (for optimization)
573    fn uses_context(&self) -> bool {
574        false
575    }
576
577    /// Get the tool's input schema
578    fn input_schema(&self) -> Value;
579}
580
581/// Handler for a tool that can complete or return an SEP-2322
582/// [`RequestOutcome::InputRequired`] continuation.
583#[cfg(feature = "stateless")]
584pub trait MrtrToolHandler: Send + Sync {
585    /// Execute an MRTR-capable tool with request context and raw arguments.
586    fn call(
587        &self,
588        ctx: RequestContext,
589        args: Value,
590    ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>>;
591
592    /// Get the tool's input schema.
593    fn input_schema(&self) -> Value;
594}
595
596/// Adapts an MRTR handler to Tower's service abstraction.
597#[cfg(feature = "stateless")]
598struct MrtrToolHandlerService<H> {
599    handler: Arc<H>,
600}
601
602#[cfg(feature = "stateless")]
603impl<H> MrtrToolHandlerService<H> {
604    fn new(handler: H) -> Self {
605        Self {
606            handler: Arc::new(handler),
607        }
608    }
609}
610
611#[cfg(feature = "stateless")]
612impl<H> Clone for MrtrToolHandlerService<H> {
613    fn clone(&self) -> Self {
614        Self {
615            handler: self.handler.clone(),
616        }
617    }
618}
619
620#[cfg(feature = "stateless")]
621impl<H> Service<ToolRequest> for MrtrToolHandlerService<H>
622where
623    H: MrtrToolHandler + 'static,
624{
625    type Response = RequestOutcome<CallToolResult>;
626    type Error = Error;
627    type Future =
628        Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
629
630    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
631        Poll::Ready(Ok(()))
632    }
633
634    fn call(&mut self, req: ToolRequest) -> Self::Future {
635        let handler = self.handler.clone();
636        Box::pin(async move { handler.call(req.ctx, req.args).await })
637    }
638}
639
640/// Runs an erased MRTR Tower service as an MRTR handler.
641#[cfg(feature = "stateless")]
642struct ServiceMrtrToolHandler {
643    service: Mutex<BoxMrtrToolService>,
644    input_schema: Value,
645}
646
647#[cfg(feature = "stateless")]
648struct GuardedMrtrToolHandler<G> {
649    guard: G,
650    inner: Arc<dyn MrtrToolHandler>,
651}
652
653#[cfg(feature = "stateless")]
654impl<G> MrtrToolHandler for GuardedMrtrToolHandler<G>
655where
656    G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
657{
658    fn call(
659        &self,
660        ctx: RequestContext,
661        args: Value,
662    ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
663        let request = ToolRequest::new(ctx, args);
664        match (self.guard)(&request) {
665            Ok(()) => self.inner.call(request.ctx, request.args),
666            Err(message) => {
667                Box::pin(
668                    async move { Ok(RequestOutcome::Complete(CallToolResult::error(message))) },
669                )
670            }
671        }
672    }
673
674    fn input_schema(&self) -> Value {
675        self.inner.input_schema()
676    }
677}
678
679#[cfg(feature = "stateless")]
680impl MrtrToolHandler for ServiceMrtrToolHandler {
681    fn call(
682        &self,
683        ctx: RequestContext,
684        args: Value,
685    ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
686        Box::pin(async move {
687            let mut service = self.service.lock().await.clone();
688            let outcome = service
689                .ready()
690                .await
691                .expect("MRTR tool service is infallible")
692                .call(ToolRequest::new(ctx, args))
693                .await
694                .expect("MRTR tool service is infallible");
695            Ok(outcome)
696        })
697    }
698
699    fn input_schema(&self) -> Value {
700        self.input_schema.clone()
701    }
702}
703
704/// Adapts a `ToolHandler` to a Tower `Service<ToolRequest>`.
705///
706/// This is an internal adapter that bridges the handler abstraction to the
707/// service abstraction, enabling middleware composition.
708pub(crate) struct ToolHandlerService<H> {
709    handler: Arc<H>,
710}
711
712impl<H> ToolHandlerService<H> {
713    pub(crate) fn new(handler: H) -> Self {
714        Self {
715            handler: Arc::new(handler),
716        }
717    }
718}
719
720impl<H> Clone for ToolHandlerService<H> {
721    fn clone(&self) -> Self {
722        Self {
723            handler: self.handler.clone(),
724        }
725    }
726}
727
728impl<H> Service<ToolRequest> for ToolHandlerService<H>
729where
730    H: ToolHandler + 'static,
731{
732    type Response = CallToolResult;
733    type Error = Error;
734    type Future = Pin<Box<dyn Future<Output = std::result::Result<CallToolResult, Error>> + Send>>;
735
736    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
737        Poll::Ready(Ok(()))
738    }
739
740    fn call(&mut self, req: ToolRequest) -> Self::Future {
741        let handler = self.handler.clone();
742        Box::pin(async move { handler.call_with_context(req.ctx, req.args).await })
743    }
744}
745
746/// A complete tool definition with service-based execution.
747///
748/// Tools are implemented as Tower services internally, enabling middleware
749/// composition via the builder's `.layer()` method. The service is wrapped
750/// in [`ToolCatchError`] to convert any errors (from handlers or middleware)
751/// into `CallToolResult::error()` responses.
752pub struct Tool {
753    /// Tool name (must be 1-128 chars, alphanumeric/underscore/hyphen/dot only)
754    pub name: String,
755    /// Human-readable title for the tool
756    pub title: Option<String>,
757    /// Description of what the tool does
758    pub description: Option<String>,
759    /// JSON Schema for the tool's output (optional)
760    pub output_schema: Option<Value>,
761    /// Icons for the tool
762    pub icons: Option<Vec<ToolIcon>>,
763    /// Tool annotations (hints about behavior)
764    pub annotations: Option<ToolAnnotations>,
765    /// Validated protocol metadata included in `tools/list`.
766    pub meta: Option<Value>,
767    /// Task support mode for this tool
768    pub task_support: TaskSupportMode,
769    /// Client capabilities required to invoke this tool in the modern
770    /// per-request protocol.
771    pub(crate) required_client_capabilities: Option<ClientCapabilities>,
772    /// Optional callback run after task allocation and before task execution.
773    pub(crate) task_preparer: Option<Arc<dyn TaskPreparer>>,
774    /// The boxed service that executes the tool
775    pub(crate) service: Option<BoxToolService>,
776    #[cfg(feature = "stateless")]
777    pub(crate) mrtr_handler: Option<Arc<dyn MrtrToolHandler>>,
778    /// JSON Schema for the tool's input
779    pub(crate) input_schema: Value,
780}
781
782impl std::fmt::Debug for Tool {
783    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
784        f.debug_struct("Tool")
785            .field("name", &self.name)
786            .field("title", &self.title)
787            .field("description", &self.description)
788            .field("output_schema", &self.output_schema)
789            .field("icons", &self.icons)
790            .field("annotations", &self.annotations)
791            .field("meta", &self.meta)
792            .field("task_support", &self.task_support)
793            .field(
794                "required_client_capabilities",
795                &self.required_client_capabilities,
796            )
797            .finish_non_exhaustive()
798    }
799}
800
801// SAFETY: BoxCloneService is Send + Sync (tower provides unsafe impl Sync),
802// and all other fields in Tool are Send + Sync.
803unsafe impl Send for Tool {}
804unsafe impl Sync for Tool {}
805
806impl Clone for Tool {
807    fn clone(&self) -> Self {
808        Self {
809            name: self.name.clone(),
810            title: self.title.clone(),
811            description: self.description.clone(),
812            output_schema: self.output_schema.clone(),
813            icons: self.icons.clone(),
814            annotations: self.annotations.clone(),
815            meta: self.meta.clone(),
816            task_support: self.task_support,
817            required_client_capabilities: self.required_client_capabilities.clone(),
818            task_preparer: self.task_preparer.clone(),
819            service: self.service.clone(),
820            #[cfg(feature = "stateless")]
821            mrtr_handler: self.mrtr_handler.clone(),
822            input_schema: self.input_schema.clone(),
823        }
824    }
825}
826
827impl Tool {
828    /// Create a new tool builder
829    pub fn builder(name: impl Into<String>) -> ToolBuilder {
830        ToolBuilder::new(name)
831    }
832
833    /// Get the tool definition for tools/list
834    pub fn definition(&self) -> ToolDefinition {
835        let execution = match self.task_support {
836            TaskSupportMode::Forbidden => None,
837            mode => Some(ToolExecution {
838                task_support: Some(mode),
839            }),
840        };
841        ToolDefinition {
842            name: self.name.clone(),
843            title: self.title.clone(),
844            description: self.description.clone(),
845            input_schema: self.input_schema.clone(),
846            output_schema: self.output_schema.clone(),
847            icons: self.icons.clone(),
848            annotations: self.annotations.clone(),
849            execution,
850            meta: self.meta.clone(),
851        }
852    }
853
854    /// Attach validated protocol metadata to this tool definition.
855    pub fn with_meta(
856        mut self,
857        meta: Value,
858    ) -> std::result::Result<Self, crate::protocol::MetaValidationError> {
859        crate::protocol::validate_meta_object(&meta)?;
860        self.meta = Some(meta);
861        Ok(self)
862    }
863
864    /// Call the tool without context
865    ///
866    /// Creates a dummy request context. For full context support, use
867    /// [`call_with_context`](Self::call_with_context).
868    pub fn call(&self, args: Value) -> BoxFuture<'static, CallToolResult> {
869        let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
870        self.call_with_context(ctx, args)
871    }
872
873    /// Call the tool with request context
874    ///
875    /// The context provides progress reporting, cancellation support, and
876    /// access to client requests (for sampling, etc.).
877    ///
878    /// # Note
879    ///
880    /// This method returns `CallToolResult` directly (not `Result<CallToolResult>`).
881    /// Any errors from the handler or middleware are converted to
882    /// `CallToolResult::error()` with `is_error: true`.
883    pub fn call_with_context(
884        &self,
885        ctx: RequestContext,
886        args: Value,
887    ) -> BoxFuture<'static, CallToolResult> {
888        let tool = self.clone();
889        Box::pin(async move {
890            match tool.call_outcome_with_context(ctx, args).await {
891                Ok(RequestOutcome::Complete(result)) => result,
892                Ok(RequestOutcome::InputRequired(_)) => CallToolResult::error(
893                    "tool requires additional client input; use call_outcome_with_context",
894                ),
895                Err(error) => CallToolResult::error(error.to_string()),
896            }
897        })
898    }
899
900    /// Call the tool and preserve an SEP-2322 input-required outcome.
901    pub fn call_outcome(
902        &self,
903        args: Value,
904    ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
905        let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
906        self.call_outcome_with_context(ctx, args)
907    }
908
909    /// Call the tool with context and preserve an SEP-2322 input-required
910    /// outcome or protocol-level handler error.
911    pub fn call_outcome_with_context(
912        &self,
913        ctx: RequestContext,
914        args: Value,
915    ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
916        use tower::ServiceExt;
917        #[cfg(feature = "stateless")]
918        if let Some(handler) = self.mrtr_handler.clone() {
919            return Box::pin(async move { handler.call(ctx, args).await });
920        }
921        let service = self
922            .service
923            .clone()
924            .expect("tool must have a complete or MRTR handler");
925        Box::pin(async move {
926            let result = service.oneshot(ToolRequest::new(ctx, args)).await.unwrap();
927            Ok(RequestOutcome::Complete(result))
928        })
929    }
930
931    /// Require the given client capability shape before this tool may be
932    /// invoked using the modern per-request protocol.
933    ///
934    /// Required objects are matched recursively. For example, requiring
935    /// `ClientCapabilities { sampling: Some(Default::default()), .. }`
936    /// accepts any advertised `sampling` capability, including one with
937    /// additional optional fields.
938    pub fn require_client_capabilities(mut self, required: ClientCapabilities) -> Self {
939        self.required_client_capabilities = Some(required);
940        self
941    }
942
943    /// Return the client capability shape required by this tool, if any.
944    pub fn required_client_capabilities(&self) -> Option<&ClientCapabilities> {
945        self.required_client_capabilities.as_ref()
946    }
947
948    /// Add a preparation callback for task-backed invocations.
949    ///
950    /// The callback runs exactly once after task ID allocation and before the
951    /// initial task response. It is skipped for synchronous calls.
952    pub fn with_task_preparation<F, Fut>(mut self, prepare: F) -> Self
953    where
954        F: Fn(TaskContext, Value) -> Fut + Send + Sync + 'static,
955        Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
956    {
957        self.task_preparer = Some(Arc::new(prepare));
958        self
959    }
960
961    /// Add a typed preparation callback to an already-built tool.
962    pub fn with_typed_task_preparation<I, F, Fut>(mut self, prepare: F) -> Self
963    where
964        I: DeserializeOwned + Send + Sync + 'static,
965        F: Fn(TaskContext, I) -> Fut + Send + Sync + 'static,
966        Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
967    {
968        self.task_preparer = Some(Arc::new(TypedTaskPreparer {
969            prepare,
970            _phantom: std::marker::PhantomData,
971        }));
972        self
973    }
974
975    pub(crate) async fn prepare_task(
976        &self,
977        context: TaskContext,
978        arguments: Value,
979    ) -> Result<TaskPreparation> {
980        match self.task_preparer.as_ref() {
981            Some(prepare) => prepare.prepare(context, arguments).await,
982            None => Ok(TaskPreparation::default()),
983        }
984    }
985
986    /// Apply a guard to this built tool.
987    ///
988    /// The guard runs before the handler and can short-circuit with an error.
989    /// This is useful for applying the same guard to multiple tools (per-group
990    /// pattern):
991    ///
992    /// ```rust
993    /// use tower_mcp::{ToolBuilder, CallToolResult};
994    /// use tower_mcp::tool::ToolRequest;
995    /// use schemars::JsonSchema;
996    /// use serde::Deserialize;
997    ///
998    /// #[derive(Debug, Deserialize, JsonSchema)]
999    /// struct Input { value: String }
1000    ///
1001    /// fn build_tool(name: &str) -> tower_mcp::tool::Tool {
1002    ///     ToolBuilder::new(name)
1003    ///         .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1004    ///         .build()
1005    /// }
1006    ///
1007    /// let guard = |_req: &ToolRequest| -> Result<(), String> { Ok(()) };
1008    ///
1009    /// let tools: Vec<_> = vec![build_tool("a"), build_tool("b")]
1010    ///     .into_iter()
1011    ///     .map(|t| t.with_guard(guard.clone()))
1012    ///     .collect();
1013    /// ```
1014    pub fn with_guard<G>(self, guard: G) -> Self
1015    where
1016        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1017    {
1018        #[cfg(feature = "stateless")]
1019        if let Some(inner) = self.mrtr_handler.clone() {
1020            return Tool {
1021                mrtr_handler: Some(Arc::new(GuardedMrtrToolHandler { guard, inner })),
1022                ..self
1023            };
1024        }
1025
1026        let guarded = GuardService {
1027            guard,
1028            inner: self
1029                .service
1030                .expect("tool must have a complete or MRTR handler"),
1031        };
1032        let caught = ToolCatchError::new(guarded);
1033        Tool {
1034            service: Some(BoxCloneService::new(caught)),
1035            ..self
1036        }
1037    }
1038
1039    /// Create a new tool with a prefixed name.
1040    ///
1041    /// This creates a copy of the tool with its name prefixed by the given
1042    /// string and a dot separator. For example, if the tool is named "query"
1043    /// and the prefix is "db", the new tool will be named "db.query".
1044    ///
1045    /// This is used internally by `McpRouter::nest()` to namespace tools.
1046    ///
1047    /// # Example
1048    ///
1049    /// ```rust
1050    /// use tower_mcp::{ToolBuilder, CallToolResult};
1051    /// use schemars::JsonSchema;
1052    /// use serde::Deserialize;
1053    ///
1054    /// #[derive(Debug, Deserialize, JsonSchema)]
1055    /// struct Input { value: String }
1056    ///
1057    /// let tool = ToolBuilder::new("query")
1058    ///     .description("Query the database")
1059    ///     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1060    ///     .build();
1061    ///
1062    /// let prefixed = tool.with_name_prefix("db");
1063    /// assert_eq!(prefixed.name, "db.query");
1064    /// ```
1065    pub fn with_name_prefix(&self, prefix: &str) -> Self {
1066        Self {
1067            name: format!("{}.{}", prefix, self.name),
1068            title: self.title.clone(),
1069            description: self.description.clone(),
1070            output_schema: self.output_schema.clone(),
1071            icons: self.icons.clone(),
1072            annotations: self.annotations.clone(),
1073            meta: self.meta.clone(),
1074            task_support: self.task_support,
1075            required_client_capabilities: self.required_client_capabilities.clone(),
1076            task_preparer: self.task_preparer.clone(),
1077            service: self.service.clone(),
1078            #[cfg(feature = "stateless")]
1079            mrtr_handler: self.mrtr_handler.clone(),
1080            input_schema: self.input_schema.clone(),
1081        }
1082    }
1083
1084    /// Create a tool from a handler (internal helper)
1085    #[allow(clippy::too_many_arguments)]
1086    fn from_handler<H: ToolHandler + 'static>(
1087        name: String,
1088        title: Option<String>,
1089        description: Option<String>,
1090        output_schema: Option<Value>,
1091        icons: Option<Vec<ToolIcon>>,
1092        annotations: Option<ToolAnnotations>,
1093        task_support: TaskSupportMode,
1094        input_schema_override: Option<Value>,
1095        handler: H,
1096    ) -> Self {
1097        let input_schema =
1098            ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
1099        let handler_service = ToolHandlerService::new(handler);
1100        let catch_error = ToolCatchError::new(handler_service);
1101        let service = BoxCloneService::new(catch_error);
1102
1103        Self {
1104            name,
1105            title,
1106            description,
1107            output_schema,
1108            icons,
1109            annotations,
1110            meta: None,
1111            task_support,
1112            required_client_capabilities: None,
1113            task_preparer: None,
1114            service: Some(service),
1115            #[cfg(feature = "stateless")]
1116            mrtr_handler: None,
1117            input_schema,
1118        }
1119    }
1120
1121    #[cfg(feature = "stateless")]
1122    #[allow(clippy::too_many_arguments)]
1123    fn from_mrtr_handler<H: MrtrToolHandler + 'static>(
1124        name: String,
1125        title: Option<String>,
1126        description: Option<String>,
1127        output_schema: Option<Value>,
1128        icons: Option<Vec<ToolIcon>>,
1129        annotations: Option<ToolAnnotations>,
1130        task_support: TaskSupportMode,
1131        input_schema_override: Option<Value>,
1132        handler: H,
1133    ) -> Self {
1134        let input_schema =
1135            ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
1136        Self {
1137            name,
1138            title,
1139            description,
1140            output_schema,
1141            icons,
1142            annotations,
1143            meta: None,
1144            task_support,
1145            required_client_capabilities: None,
1146            task_preparer: None,
1147            service: None,
1148            mrtr_handler: Some(Arc::new(handler)),
1149            input_schema,
1150        }
1151    }
1152}
1153
1154// =============================================================================
1155// Builder API
1156// =============================================================================
1157
1158/// Builder for creating tools with a fluent API
1159///
1160/// # Example
1161///
1162/// ```rust
1163/// use tower_mcp::{ToolBuilder, CallToolResult};
1164/// use schemars::JsonSchema;
1165/// use serde::Deserialize;
1166///
1167/// #[derive(Debug, Deserialize, JsonSchema)]
1168/// struct GreetInput {
1169///     name: String,
1170/// }
1171///
1172/// let tool = ToolBuilder::new("greet")
1173///     .description("Greet someone by name")
1174///     .handler(|input: GreetInput| async move {
1175///         Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
1176///     })
1177///     .build();
1178///
1179/// assert_eq!(tool.name, "greet");
1180/// ```
1181pub struct ToolBuilder {
1182    name: String,
1183    title: Option<String>,
1184    description: Option<String>,
1185    output_schema: Option<Value>,
1186    input_schema_override: Option<Value>,
1187    icons: Option<Vec<ToolIcon>>,
1188    annotations: Option<ToolAnnotations>,
1189    task_support: TaskSupportMode,
1190}
1191
1192impl ToolBuilder {
1193    /// Create a new tool builder with the given name.
1194    ///
1195    /// Tool names must be 1-64 characters and contain only ASCII alphanumeric
1196    /// characters, underscores, hyphens, dots, and forward slashes (per
1197    /// [SEP-986](https://github.com/modelcontextprotocol/specification/issues/986)).
1198    ///
1199    /// Use [`try_new`](Self::try_new) if the name comes from runtime input.
1200    ///
1201    /// # Panics
1202    ///
1203    /// Panics if `name` is empty, exceeds 64 characters, or contains
1204    /// characters other than ASCII alphanumerics, `_`, `-`, `.`, and `/`.
1205    pub fn new(name: impl Into<String>) -> Self {
1206        let name = name.into();
1207        if let Err(e) = validate_tool_name(&name) {
1208            panic!("{e}");
1209        }
1210        Self {
1211            name,
1212            title: None,
1213            description: None,
1214            output_schema: None,
1215            input_schema_override: None,
1216            icons: None,
1217            annotations: None,
1218            task_support: TaskSupportMode::default(),
1219        }
1220    }
1221
1222    /// Create a new tool builder, returning an error if the name is invalid.
1223    ///
1224    /// This is the fallible alternative to [`new`](Self::new) for cases where
1225    /// the tool name comes from runtime input (e.g., user configuration or
1226    /// database).
1227    pub fn try_new(name: impl Into<String>) -> Result<Self> {
1228        let name = name.into();
1229        validate_tool_name(&name)?;
1230        Ok(Self {
1231            name,
1232            title: None,
1233            description: None,
1234            output_schema: None,
1235            input_schema_override: None,
1236            icons: None,
1237            annotations: None,
1238            task_support: TaskSupportMode::default(),
1239        })
1240    }
1241
1242    /// Set a human-readable title for the tool.
1243    ///
1244    /// The title is displayed by MCP clients (e.g., Claude Code's `/mcp` tool list)
1245    /// as a friendly label instead of the raw tool name. For example, a tool named
1246    /// `search_crates` with title `"Search Crates"` will display the title in UIs
1247    /// that support it.
1248    ///
1249    /// ```
1250    /// # use tower_mcp::ToolBuilder;
1251    /// let tool = ToolBuilder::new("search_crates")
1252    ///     .title("Search Crates")
1253    ///     .description("Search for Rust crates on crates.io")
1254    ///     .handler(|()| async { Ok(tower_mcp::CallToolResult::text("results")) })
1255    ///     .build();
1256    /// ```
1257    pub fn title(mut self, title: impl Into<String>) -> Self {
1258        self.title = Some(title.into());
1259        self
1260    }
1261
1262    /// Set the output schema (JSON Schema for structured output)
1263    pub fn output_schema(mut self, schema: Value) -> Self {
1264        self.output_schema = Some(schema);
1265        self
1266    }
1267
1268    /// Override the input schema (JSON Schema for tool arguments).
1269    ///
1270    /// By default, the input schema is auto-generated from the handler's input
1271    /// type via [`schemars::JsonSchema`]. Calling this method overrides that
1272    /// auto-generation with an explicit schema. This is particularly useful for
1273    /// handlers that use [`RawArgs`](crate::extract::RawArgs) (which has no typed
1274    /// input struct) but still need to declare a non-trivial schema, or to
1275    /// supply richer JSON Schema 2020-12 constructs (`oneOf`, `anyOf`,
1276    /// `if`/`then`, `$ref`, etc.) that schemars cannot express.
1277    ///
1278    /// The supplied schema is normalized via the same `type: "object"` check
1279    /// the auto-generated schemas go through, so MCP-spec compliance is
1280    /// preserved.
1281    ///
1282    /// When called alongside a typed handler (`.handler(|x: Foo| ...)` or a
1283    /// [`Json<T>`](crate::extract::Json) extractor), the explicit schema wins
1284    /// over the schemars-generated one.
1285    ///
1286    /// # Example
1287    ///
1288    /// ```rust
1289    /// use serde_json::json;
1290    /// use tower_mcp::{CallToolResult, ToolBuilder};
1291    /// use tower_mcp::extract::RawArgs;
1292    ///
1293    /// let tool = ToolBuilder::new("query")
1294    ///     .description("Query with a conditional schema")
1295    ///     .input_schema(json!({
1296    ///         "type": "object",
1297    ///         "properties": {
1298    ///             "filter": {
1299    ///                 "oneOf": [
1300    ///                     { "type": "string" },
1301    ///                     {
1302    ///                         "type": "object",
1303    ///                         "properties": { "field": { "type": "string" } },
1304    ///                         "required": ["field"]
1305    ///                     }
1306    ///                 ]
1307    ///             }
1308    ///         },
1309    ///         "required": ["filter"]
1310    ///     }))
1311    ///     .extractor_handler((), |RawArgs(args): RawArgs| async move {
1312    ///         Ok(CallToolResult::json(args))
1313    ///     })
1314    ///     .build();
1315    ///
1316    /// let schema = tool.definition().input_schema;
1317    /// assert_eq!(schema["type"], "object");
1318    /// assert!(schema["properties"]["filter"]["oneOf"].is_array());
1319    /// ```
1320    pub fn input_schema(mut self, schema: Value) -> Self {
1321        self.input_schema_override = Some(schema);
1322        self
1323    }
1324
1325    /// Add an icon for the tool
1326    pub fn icon(mut self, src: impl Into<String>) -> Self {
1327        self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1328            src: src.into(),
1329            mime_type: None,
1330            sizes: None,
1331            theme: None,
1332        });
1333        self
1334    }
1335
1336    /// Add an icon with metadata
1337    pub fn icon_with_meta(
1338        mut self,
1339        src: impl Into<String>,
1340        mime_type: Option<String>,
1341        sizes: Option<Vec<String>>,
1342    ) -> Self {
1343        self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1344            src: src.into(),
1345            mime_type,
1346            sizes,
1347            theme: None,
1348        });
1349        self
1350    }
1351
1352    /// Set the tool description
1353    pub fn description(mut self, description: impl Into<String>) -> Self {
1354        self.description = Some(description.into());
1355        self
1356    }
1357
1358    /// Mark the tool as read-only (does not modify state)
1359    pub fn read_only(mut self) -> Self {
1360        self.annotations
1361            .get_or_insert_with(ToolAnnotations::default)
1362            .read_only_hint = true;
1363        self
1364    }
1365
1366    /// Mark the tool as non-destructive
1367    pub fn non_destructive(mut self) -> Self {
1368        self.annotations
1369            .get_or_insert_with(ToolAnnotations::default)
1370            .destructive_hint = false;
1371        self
1372    }
1373
1374    /// Mark the tool as destructive (may perform irreversible operations)
1375    pub fn destructive(mut self) -> Self {
1376        self.annotations
1377            .get_or_insert_with(ToolAnnotations::default)
1378            .destructive_hint = true;
1379        self
1380    }
1381
1382    /// Mark the tool as idempotent (same args = same effect)
1383    pub fn idempotent(mut self) -> Self {
1384        self.annotations
1385            .get_or_insert_with(ToolAnnotations::default)
1386            .idempotent_hint = true;
1387        self
1388    }
1389
1390    /// Mark the tool as read-only, idempotent, and non-destructive.
1391    ///
1392    /// This is a convenience method for safe, side-effect-free tools.
1393    /// For finer control, use `.read_only()`, `.idempotent()`, and
1394    /// `.non_destructive()` individually.
1395    pub fn read_only_safe(mut self) -> Self {
1396        let ann = self
1397            .annotations
1398            .get_or_insert_with(ToolAnnotations::default);
1399        ann.read_only_hint = true;
1400        ann.idempotent_hint = true;
1401        ann.destructive_hint = false;
1402        self
1403    }
1404
1405    /// Set tool annotations directly
1406    pub fn annotations(mut self, annotations: ToolAnnotations) -> Self {
1407        self.annotations = Some(annotations);
1408        self
1409    }
1410
1411    /// Set the task support mode for this tool
1412    pub fn task_support(mut self, mode: TaskSupportMode) -> Self {
1413        self.task_support = mode;
1414        self
1415    }
1416
1417    /// Create a tool that takes no parameters.
1418    ///
1419    /// This is a convenience method for tools that don't require any input.
1420    /// It generates the correct `{"type": "object"}` schema that MCP clients expect.
1421    ///
1422    /// # Example
1423    ///
1424    /// ```rust
1425    /// use tower_mcp::{ToolBuilder, CallToolResult};
1426    ///
1427    /// let tool = ToolBuilder::new("get_status")
1428    ///     .description("Get current status")
1429    ///     .no_params_handler(|| async {
1430    ///         Ok(CallToolResult::text("OK"))
1431    ///     })
1432    ///     .build();
1433    /// ```
1434    pub fn no_params_handler<F, Fut>(self, handler: F) -> ToolBuilderWithNoParamsHandler<F>
1435    where
1436        F: Fn() -> Fut + Send + Sync + 'static,
1437        Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1438    {
1439        ToolBuilderWithNoParamsHandler {
1440            name: self.name,
1441            title: self.title,
1442            description: self.description,
1443            output_schema: self.output_schema,
1444            input_schema_override: self.input_schema_override,
1445            icons: self.icons,
1446            annotations: self.annotations,
1447            task_support: self.task_support,
1448            handler,
1449        }
1450    }
1451
1452    /// Specify input type and handler.
1453    ///
1454    /// The input type must implement `JsonSchema` and `DeserializeOwned`.
1455    /// The handler receives the deserialized input and returns a `CallToolResult`.
1456    ///
1457    /// # State Sharing
1458    ///
1459    /// To share state across tool calls (e.g., database connections, API clients),
1460    /// wrap your state in an `Arc` and clone it into the async block:
1461    ///
1462    /// ```rust
1463    /// use std::sync::Arc;
1464    /// use tower_mcp::{ToolBuilder, CallToolResult};
1465    /// use schemars::JsonSchema;
1466    /// use serde::Deserialize;
1467    ///
1468    /// struct AppState {
1469    ///     api_key: String,
1470    /// }
1471    ///
1472    /// #[derive(Debug, Deserialize, JsonSchema)]
1473    /// struct MyInput {
1474    ///     query: String,
1475    /// }
1476    ///
1477    /// let state = Arc::new(AppState { api_key: "secret".to_string() });
1478    ///
1479    /// let tool = ToolBuilder::new("my_tool")
1480    ///     .description("A tool that uses shared state")
1481    ///     .handler(move |input: MyInput| {
1482    ///         let state = state.clone(); // Clone Arc for the async block
1483    ///         async move {
1484    ///             // Use state.api_key here...
1485    ///             Ok(CallToolResult::text(format!("Query: {}", input.query)))
1486    ///         }
1487    ///     })
1488    ///     .build();
1489    /// ```
1490    ///
1491    /// The `move` keyword on the closure captures the `Arc<AppState>`, and
1492    /// cloning it inside the closure body allows each async invocation to
1493    /// have its own reference to the shared state.
1494    pub fn handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithHandler<I, F>
1495    where
1496        I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1497        F: Fn(I) -> Fut + Send + Sync + 'static,
1498        Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1499    {
1500        ToolBuilderWithHandler {
1501            name: self.name,
1502            title: self.title,
1503            description: self.description,
1504            output_schema: self.output_schema,
1505            input_schema_override: self.input_schema_override,
1506            icons: self.icons,
1507            annotations: self.annotations,
1508            task_support: self.task_support,
1509            task_preparer: None,
1510            handler,
1511            _phantom: std::marker::PhantomData,
1512        }
1513    }
1514
1515    /// Set an SEP-2322 handler that may return either a complete tool result
1516    /// or an input-required continuation.
1517    ///
1518    /// The handler receives [`RequestContext`], where
1519    /// [`RequestContext::input_responses`] and
1520    /// [`RequestContext::request_state`] expose values from a retry.
1521    #[cfg(feature = "stateless")]
1522    pub fn mrtr_handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithMrtrHandler<I, F>
1523    where
1524        I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1525        F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
1526        Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
1527    {
1528        ToolBuilderWithMrtrHandler {
1529            name: self.name,
1530            title: self.title,
1531            description: self.description,
1532            output_schema: self.output_schema,
1533            input_schema_override: self.input_schema_override,
1534            icons: self.icons,
1535            annotations: self.annotations,
1536            task_support: self.task_support,
1537            handler,
1538            _phantom: std::marker::PhantomData,
1539        }
1540    }
1541
1542    /// Create a tool using the extractor pattern.
1543    ///
1544    /// This method provides an axum-inspired way to define handlers where state,
1545    /// context, and input are extracted declaratively from function parameters.
1546    /// This reduces the combinatorial explosion of handler variants like
1547    /// `handler_with_state`, `handler_with_context`, etc.
1548    ///
1549    /// # Schema Auto-Detection
1550    ///
1551    /// When a [`Json<T>`](crate::extract::Json) extractor is used, the proper JSON
1552    /// schema is automatically generated from `T`'s `JsonSchema` implementation.
1553    /// No turbofish is needed -- the schema type is inferred from the closure
1554    /// parameters.
1555    ///
1556    /// # Extractors
1557    ///
1558    /// Built-in extractors available in [`crate::extract`]:
1559    /// - [`Json<T>`](crate::extract::Json) - Deserialize JSON arguments to type `T`
1560    /// - [`State<T>`](crate::extract::State) - Extract cloned state
1561    /// - [`Extension<T>`](crate::extract::Extension) - Extract router-level state
1562    /// - [`Context`](crate::extract::Context) - Extract request context
1563    /// - [`RawArgs`](crate::extract::RawArgs) - Extract raw JSON arguments
1564    ///
1565    /// # Per-Tool Middleware
1566    ///
1567    /// The returned builder supports `.layer()` to apply Tower middleware:
1568    ///
1569    /// ```rust
1570    /// use std::sync::Arc;
1571    /// use std::time::Duration;
1572    /// use tower::timeout::TimeoutLayer;
1573    /// use tower_mcp::{ToolBuilder, CallToolResult};
1574    /// use tower_mcp::extract::{Json, State};
1575    /// use schemars::JsonSchema;
1576    /// use serde::Deserialize;
1577    ///
1578    /// #[derive(Clone)]
1579    /// struct Database { url: String }
1580    ///
1581    /// #[derive(Debug, Deserialize, JsonSchema)]
1582    /// struct QueryInput { query: String }
1583    ///
1584    /// let db = Arc::new(Database { url: "postgres://...".to_string() });
1585    ///
1586    /// let tool = ToolBuilder::new("search")
1587    ///     .description("Search the database")
1588    ///     .extractor_handler(db, |
1589    ///         State(db): State<Arc<Database>>,
1590    ///         Json(input): Json<QueryInput>,
1591    ///     | async move {
1592    ///         Ok(CallToolResult::text(format!("Searched {} with: {}", db.url, input.query)))
1593    ///     })
1594    ///     .layer(TimeoutLayer::new(Duration::from_secs(30)))
1595    ///     .build();
1596    /// ```
1597    ///
1598    /// # Example
1599    ///
1600    /// ```rust
1601    /// use std::sync::Arc;
1602    /// use tower_mcp::{ToolBuilder, CallToolResult};
1603    /// use tower_mcp::extract::{Json, State, Context};
1604    /// use schemars::JsonSchema;
1605    /// use serde::Deserialize;
1606    ///
1607    /// #[derive(Clone)]
1608    /// struct Database { url: String }
1609    ///
1610    /// #[derive(Debug, Deserialize, JsonSchema)]
1611    /// struct QueryInput { query: String }
1612    ///
1613    /// let db = Arc::new(Database { url: "postgres://...".to_string() });
1614    ///
1615    /// let tool = ToolBuilder::new("search")
1616    ///     .description("Search the database")
1617    ///     .extractor_handler(db, |
1618    ///         State(db): State<Arc<Database>>,
1619    ///         ctx: Context,
1620    ///         Json(input): Json<QueryInput>,
1621    ///     | async move {
1622    ///         if ctx.is_cancelled() {
1623    ///             return Ok(CallToolResult::error("Cancelled"));
1624    ///         }
1625    ///         ctx.report_progress(0.5, Some(1.0), Some("Searching...")).await;
1626    ///         Ok(CallToolResult::text(format!("Searched {} with: {}", db.url, input.query)))
1627    ///     })
1628    ///     .build();
1629    /// ```
1630    ///
1631    /// # Type Inference
1632    ///
1633    /// The compiler infers extractor types from the function signature. Make sure
1634    /// to annotate the extractor types explicitly in the closure parameters.
1635    pub fn extractor_handler<S, F, T>(
1636        self,
1637        state: S,
1638        handler: F,
1639    ) -> crate::extract::ToolBuilderWithExtractor<S, F, T>
1640    where
1641        S: Clone + Send + Sync + 'static,
1642        F: crate::extract::ExtractorHandler<S, T> + Clone,
1643        T: Send + Sync + 'static,
1644    {
1645        let input_schema = ensure_object_schema(
1646            self.input_schema_override
1647                .unwrap_or_else(|| F::input_schema()),
1648        );
1649        crate::extract::ToolBuilderWithExtractor {
1650            name: self.name,
1651            title: self.title,
1652            description: self.description,
1653            output_schema: self.output_schema,
1654            icons: self.icons,
1655            annotations: self.annotations,
1656            task_support: self.task_support,
1657            state,
1658            handler,
1659            input_schema,
1660            _phantom: std::marker::PhantomData,
1661        }
1662    }
1663
1664    /// Create a tool using the extractor pattern with typed JSON input.
1665    ///
1666    /// # Deprecated
1667    ///
1668    /// Use [`extractor_handler`](Self::extractor_handler) instead. It auto-detects
1669    /// the JSON schema from `Json<T>` extractors, producing identical results
1670    /// without requiring a turbofish.
1671    ///
1672    /// ```rust
1673    /// # use std::sync::Arc;
1674    /// # use tower_mcp::{ToolBuilder, CallToolResult};
1675    /// # use tower_mcp::extract::{Json, State};
1676    /// # use schemars::JsonSchema;
1677    /// # use serde::Deserialize;
1678    /// # #[derive(Clone)]
1679    /// # struct AppState { prefix: String }
1680    /// # #[derive(Debug, Deserialize, JsonSchema)]
1681    /// # struct GreetInput { name: String }
1682    /// # let state = Arc::new(AppState { prefix: "Hello".to_string() });
1683    /// // Before (deprecated):
1684    /// // .extractor_handler_typed::<_, _, _, GreetInput>(state, handler)
1685    ///
1686    /// // After:
1687    /// let tool = ToolBuilder::new("greet")
1688    ///     .description("Greet someone")
1689    ///     .extractor_handler(state, |
1690    ///         State(app): State<Arc<AppState>>,
1691    ///         Json(input): Json<GreetInput>,
1692    ///     | async move {
1693    ///         Ok(CallToolResult::text(format!("{}, {}!", app.prefix, input.name)))
1694    ///     })
1695    ///     .build();
1696    /// ```
1697    #[deprecated(
1698        since = "0.8.0",
1699        note = "Use `extractor_handler` instead -- it auto-detects JSON schema from `Json<T>` extractors without requiring a turbofish"
1700    )]
1701    #[allow(deprecated)]
1702    pub fn extractor_handler_typed<S, F, T, I>(
1703        self,
1704        state: S,
1705        handler: F,
1706    ) -> crate::extract::ToolBuilderWithTypedExtractor<S, F, T, I>
1707    where
1708        S: Clone + Send + Sync + 'static,
1709        F: crate::extract::TypedExtractorHandler<S, T, I> + Clone,
1710        T: Send + Sync + 'static,
1711        I: schemars::JsonSchema + Send + Sync + 'static,
1712    {
1713        crate::extract::ToolBuilderWithTypedExtractor {
1714            name: self.name,
1715            title: self.title,
1716            description: self.description,
1717            output_schema: self.output_schema,
1718            input_schema_override: self.input_schema_override,
1719            icons: self.icons,
1720            annotations: self.annotations,
1721            task_support: self.task_support,
1722            state,
1723            handler,
1724            _phantom: std::marker::PhantomData,
1725        }
1726    }
1727}
1728
1729/// Handler for tools with no parameters.
1730///
1731/// Used internally by [`ToolBuilder::no_params_handler`].
1732struct NoParamsTypedHandler<F> {
1733    handler: F,
1734}
1735
1736impl<F, Fut> ToolHandler for NoParamsTypedHandler<F>
1737where
1738    F: Fn() -> Fut + Send + Sync + 'static,
1739    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1740{
1741    fn call(&self, _args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
1742        Box::pin(async move { (self.handler)().await })
1743    }
1744
1745    fn input_schema(&self) -> Value {
1746        serde_json::json!({ "type": "object" })
1747    }
1748}
1749
1750/// Builder state after handler is specified
1751#[doc(hidden)]
1752pub struct ToolBuilderWithHandler<I, F> {
1753    name: String,
1754    title: Option<String>,
1755    description: Option<String>,
1756    output_schema: Option<Value>,
1757    input_schema_override: Option<Value>,
1758    icons: Option<Vec<ToolIcon>>,
1759    annotations: Option<ToolAnnotations>,
1760    task_support: TaskSupportMode,
1761    task_preparer: Option<Arc<dyn TaskPreparer>>,
1762    handler: F,
1763    _phantom: std::marker::PhantomData<I>,
1764}
1765
1766/// Builder state for an SEP-2322-capable tool handler.
1767#[cfg(feature = "stateless")]
1768#[doc(hidden)]
1769pub struct ToolBuilderWithMrtrHandler<I, F> {
1770    name: String,
1771    title: Option<String>,
1772    description: Option<String>,
1773    output_schema: Option<Value>,
1774    input_schema_override: Option<Value>,
1775    icons: Option<Vec<ToolIcon>>,
1776    annotations: Option<ToolAnnotations>,
1777    task_support: TaskSupportMode,
1778    handler: F,
1779    _phantom: std::marker::PhantomData<I>,
1780}
1781
1782/// Builder state after a layer has been applied to an MRTR handler.
1783#[cfg(feature = "stateless")]
1784#[doc(hidden)]
1785pub struct ToolBuilderWithMrtrLayer<I, F, L> {
1786    name: String,
1787    title: Option<String>,
1788    description: Option<String>,
1789    output_schema: Option<Value>,
1790    input_schema_override: Option<Value>,
1791    icons: Option<Vec<ToolIcon>>,
1792    annotations: Option<ToolAnnotations>,
1793    task_support: TaskSupportMode,
1794    handler: F,
1795    layer: L,
1796    _phantom: std::marker::PhantomData<I>,
1797}
1798
1799/// Builder state for tools with no parameters.
1800///
1801/// Created by [`ToolBuilder::no_params_handler`].
1802#[doc(hidden)]
1803pub struct ToolBuilderWithNoParamsHandler<F> {
1804    name: String,
1805    title: Option<String>,
1806    description: Option<String>,
1807    output_schema: Option<Value>,
1808    input_schema_override: Option<Value>,
1809    icons: Option<Vec<ToolIcon>>,
1810    annotations: Option<ToolAnnotations>,
1811    task_support: TaskSupportMode,
1812    handler: F,
1813}
1814
1815impl<F, Fut> ToolBuilderWithNoParamsHandler<F>
1816where
1817    F: Fn() -> Fut + Send + Sync + 'static,
1818    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1819{
1820    /// Build the tool.
1821    pub fn build(self) -> Tool {
1822        Tool::from_handler(
1823            self.name,
1824            self.title,
1825            self.description,
1826            self.output_schema,
1827            self.icons,
1828            self.annotations,
1829            self.task_support,
1830            self.input_schema_override,
1831            NoParamsTypedHandler {
1832                handler: self.handler,
1833            },
1834        )
1835    }
1836
1837    /// Apply a Tower layer (middleware) to this tool.
1838    ///
1839    /// See [`ToolBuilderWithHandler::layer`] for details.
1840    pub fn layer<L>(self, layer: L) -> ToolBuilderWithNoParamsHandlerLayer<F, L> {
1841        ToolBuilderWithNoParamsHandlerLayer {
1842            name: self.name,
1843            title: self.title,
1844            description: self.description,
1845            output_schema: self.output_schema,
1846            input_schema_override: self.input_schema_override,
1847            icons: self.icons,
1848            annotations: self.annotations,
1849            task_support: self.task_support,
1850            handler: self.handler,
1851            layer,
1852        }
1853    }
1854
1855    /// Apply a guard to this tool.
1856    ///
1857    /// See [`ToolBuilderWithHandler::guard`] for details.
1858    pub fn guard<G>(self, guard: G) -> ToolBuilderWithNoParamsHandlerLayer<F, GuardLayer<G>>
1859    where
1860        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1861    {
1862        self.layer(GuardLayer::new(guard))
1863    }
1864}
1865
1866/// Builder state after a layer has been applied to a no-params handler.
1867#[doc(hidden)]
1868pub struct ToolBuilderWithNoParamsHandlerLayer<F, L> {
1869    name: String,
1870    title: Option<String>,
1871    description: Option<String>,
1872    output_schema: Option<Value>,
1873    input_schema_override: Option<Value>,
1874    icons: Option<Vec<ToolIcon>>,
1875    annotations: Option<ToolAnnotations>,
1876    task_support: TaskSupportMode,
1877    handler: F,
1878    layer: L,
1879}
1880
1881#[allow(private_bounds)]
1882impl<F, Fut, L> ToolBuilderWithNoParamsHandlerLayer<F, L>
1883where
1884    F: Fn() -> Fut + Send + Sync + 'static,
1885    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1886    L: tower::Layer<ToolHandlerService<NoParamsTypedHandler<F>>> + Clone + Send + Sync + 'static,
1887    L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
1888    <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
1889    <L::Service as Service<ToolRequest>>::Future: Send,
1890{
1891    /// Build the tool with the applied layer(s).
1892    pub fn build(self) -> Tool {
1893        let input_schema = ensure_object_schema(
1894            self.input_schema_override
1895                .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
1896        );
1897
1898        let handler_service = ToolHandlerService::new(NoParamsTypedHandler {
1899            handler: self.handler,
1900        });
1901        let layered = self.layer.layer(handler_service);
1902        let catch_error = ToolCatchError::new(layered);
1903        let service = BoxCloneService::new(catch_error);
1904
1905        Tool {
1906            name: self.name,
1907            title: self.title,
1908            description: self.description,
1909            output_schema: self.output_schema,
1910            icons: self.icons,
1911            annotations: self.annotations,
1912            meta: None,
1913            task_support: self.task_support,
1914            required_client_capabilities: None,
1915            task_preparer: None,
1916            service: Some(service),
1917            #[cfg(feature = "stateless")]
1918            mrtr_handler: None,
1919            input_schema,
1920        }
1921    }
1922
1923    /// Apply an additional Tower layer (middleware).
1924    pub fn layer<L2>(
1925        self,
1926        layer: L2,
1927    ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<L2, L>> {
1928        ToolBuilderWithNoParamsHandlerLayer {
1929            name: self.name,
1930            title: self.title,
1931            description: self.description,
1932            output_schema: self.output_schema,
1933            input_schema_override: self.input_schema_override,
1934            icons: self.icons,
1935            annotations: self.annotations,
1936            task_support: self.task_support,
1937            handler: self.handler,
1938            layer: tower::layer::util::Stack::new(layer, self.layer),
1939        }
1940    }
1941
1942    /// Apply a guard to this tool.
1943    ///
1944    /// See [`ToolBuilderWithHandler::guard`] for details.
1945    pub fn guard<G>(
1946        self,
1947        guard: G,
1948    ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<GuardLayer<G>, L>>
1949    where
1950        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1951    {
1952        self.layer(GuardLayer::new(guard))
1953    }
1954}
1955
1956impl<I, F, Fut> ToolBuilderWithHandler<I, F>
1957where
1958    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1959    F: Fn(I) -> Fut + Send + Sync + 'static,
1960    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1961{
1962    /// Build the tool.
1963    pub fn build(self) -> Tool {
1964        let mut tool = Tool::from_handler(
1965            self.name,
1966            self.title,
1967            self.description,
1968            self.output_schema,
1969            self.icons,
1970            self.annotations,
1971            self.task_support,
1972            self.input_schema_override,
1973            TypedHandler {
1974                handler: self.handler,
1975                _phantom: std::marker::PhantomData,
1976            },
1977        );
1978        tool.task_preparer = self.task_preparer;
1979        tool
1980    }
1981
1982    /// Add a typed preparation step for task-backed invocations.
1983    pub fn task_preparation<P, PrepareFuture>(mut self, prepare: P) -> Self
1984    where
1985        P: Fn(TaskContext, I) -> PrepareFuture + Send + Sync + 'static,
1986        PrepareFuture: Future<Output = Result<TaskPreparation>> + Send + 'static,
1987    {
1988        self.task_preparer = Some(Arc::new(TypedTaskPreparer {
1989            prepare,
1990            _phantom: std::marker::PhantomData,
1991        }));
1992        self
1993    }
1994
1995    /// Apply a Tower layer (middleware) to this tool.
1996    ///
1997    /// The layer wraps the tool's handler service, enabling functionality like
1998    /// timeouts, rate limiting, and metrics collection at the per-tool level.
1999    ///
2000    /// # Example
2001    ///
2002    /// ```rust
2003    /// use std::time::Duration;
2004    /// use tower::timeout::TimeoutLayer;
2005    /// use tower_mcp::{ToolBuilder, CallToolResult};
2006    /// use schemars::JsonSchema;
2007    /// use serde::Deserialize;
2008    ///
2009    /// #[derive(Debug, Deserialize, JsonSchema)]
2010    /// struct Input { query: String }
2011    ///
2012    /// let tool = ToolBuilder::new("search")
2013    ///     .description("Search with timeout")
2014    ///     .handler(|input: Input| async move {
2015    ///         Ok(CallToolResult::text("result"))
2016    ///     })
2017    ///     .layer(TimeoutLayer::new(Duration::from_secs(30)))
2018    ///     .build();
2019    /// ```
2020    pub fn layer<L>(self, layer: L) -> ToolBuilderWithLayer<I, F, L> {
2021        ToolBuilderWithLayer {
2022            name: self.name,
2023            title: self.title,
2024            description: self.description,
2025            output_schema: self.output_schema,
2026            input_schema_override: self.input_schema_override,
2027            icons: self.icons,
2028            annotations: self.annotations,
2029            task_support: self.task_support,
2030            task_preparer: self.task_preparer,
2031            handler: self.handler,
2032            layer,
2033            _phantom: std::marker::PhantomData,
2034        }
2035    }
2036
2037    /// Apply a guard to this tool.
2038    ///
2039    /// The guard runs before the handler and can short-circuit with an error
2040    /// message. This is syntactic sugar for `.layer(GuardLayer::new(f))`.
2041    ///
2042    /// See [`GuardLayer`] for a full example.
2043    pub fn guard<G>(self, guard: G) -> ToolBuilderWithLayer<I, F, GuardLayer<G>>
2044    where
2045        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2046    {
2047        self.layer(GuardLayer::new(guard))
2048    }
2049}
2050
2051#[cfg(feature = "stateless")]
2052impl<I, F, Fut> ToolBuilderWithMrtrHandler<I, F>
2053where
2054    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2055    F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2056    Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2057{
2058    /// Build the MRTR-capable tool.
2059    pub fn build(self) -> Tool {
2060        Tool::from_mrtr_handler(
2061            self.name,
2062            self.title,
2063            self.description,
2064            self.output_schema,
2065            self.icons,
2066            self.annotations,
2067            self.task_support,
2068            self.input_schema_override,
2069            TypedMrtrHandler {
2070                handler: self.handler,
2071                _phantom: std::marker::PhantomData,
2072            },
2073        )
2074    }
2075
2076    /// Apply a Tower layer to every attempt at this MRTR-capable tool.
2077    ///
2078    /// Each MRTR retry is an independent request, so the layer runs once per
2079    /// round. Middleware failures become complete tool error results, matching
2080    /// the behavior of layers on non-MRTR tools.
2081    pub fn layer<L>(self, layer: L) -> ToolBuilderWithMrtrLayer<I, F, L> {
2082        ToolBuilderWithMrtrLayer {
2083            name: self.name,
2084            title: self.title,
2085            description: self.description,
2086            output_schema: self.output_schema,
2087            input_schema_override: self.input_schema_override,
2088            icons: self.icons,
2089            annotations: self.annotations,
2090            task_support: self.task_support,
2091            handler: self.handler,
2092            layer,
2093            _phantom: std::marker::PhantomData,
2094        }
2095    }
2096
2097    /// Apply a guard to every attempt at this MRTR-capable tool.
2098    pub fn guard<G>(self, guard: G) -> ToolBuilderWithMrtrLayer<I, F, GuardLayer<G>>
2099    where
2100        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2101    {
2102        self.layer(GuardLayer::new(guard))
2103    }
2104}
2105
2106#[cfg(feature = "stateless")]
2107#[allow(private_bounds)]
2108impl<I, F, Fut, L> ToolBuilderWithMrtrLayer<I, F, L>
2109where
2110    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2111    F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2112    Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2113    L: tower::Layer<MrtrToolHandlerService<TypedMrtrHandler<I, F>>> + Clone + Send + Sync + 'static,
2114    L::Service:
2115        Service<ToolRequest, Response = RequestOutcome<CallToolResult>> + Clone + Send + 'static,
2116    <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send + 'static,
2117    <L::Service as Service<ToolRequest>>::Future: Send + 'static,
2118{
2119    /// Build the MRTR-capable tool with the applied layer(s).
2120    pub fn build(self) -> Tool {
2121        let input_schema = self.input_schema_override.unwrap_or_else(|| {
2122            let schema = schemars::schema_for!(I);
2123            serde_json::to_value(schema).unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
2124        });
2125        let input_schema = ensure_object_schema(input_schema);
2126        let service = MrtrToolHandlerService::new(TypedMrtrHandler {
2127            handler: self.handler,
2128            _phantom: std::marker::PhantomData,
2129        });
2130        let service = self.layer.layer(service);
2131        let service = BoxCloneService::new(MrtrToolCatchError::new(service));
2132
2133        Tool {
2134            name: self.name,
2135            title: self.title,
2136            description: self.description,
2137            output_schema: self.output_schema,
2138            icons: self.icons,
2139            annotations: self.annotations,
2140            meta: None,
2141            task_support: self.task_support,
2142            required_client_capabilities: None,
2143            task_preparer: None,
2144            service: None,
2145            mrtr_handler: Some(Arc::new(ServiceMrtrToolHandler {
2146                service: Mutex::new(service),
2147                input_schema: input_schema.clone(),
2148            })),
2149            input_schema,
2150        }
2151    }
2152
2153    /// Apply an additional Tower layer.
2154    pub fn layer<L2>(
2155        self,
2156        layer: L2,
2157    ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<L2, L>> {
2158        ToolBuilderWithMrtrLayer {
2159            name: self.name,
2160            title: self.title,
2161            description: self.description,
2162            output_schema: self.output_schema,
2163            input_schema_override: self.input_schema_override,
2164            icons: self.icons,
2165            annotations: self.annotations,
2166            task_support: self.task_support,
2167            handler: self.handler,
2168            layer: tower::layer::util::Stack::new(layer, self.layer),
2169            _phantom: std::marker::PhantomData,
2170        }
2171    }
2172
2173    /// Apply an additional guard.
2174    pub fn guard<G>(
2175        self,
2176        guard: G,
2177    ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2178    where
2179        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2180    {
2181        self.layer(GuardLayer::new(guard))
2182    }
2183}
2184
2185/// Builder state after a layer has been applied to the handler.
2186///
2187/// This builder allows chaining additional layers and building the final tool.
2188#[doc(hidden)]
2189pub struct ToolBuilderWithLayer<I, F, L> {
2190    name: String,
2191    title: Option<String>,
2192    description: Option<String>,
2193    output_schema: Option<Value>,
2194    input_schema_override: Option<Value>,
2195    icons: Option<Vec<ToolIcon>>,
2196    annotations: Option<ToolAnnotations>,
2197    task_support: TaskSupportMode,
2198    task_preparer: Option<Arc<dyn TaskPreparer>>,
2199    handler: F,
2200    layer: L,
2201    _phantom: std::marker::PhantomData<I>,
2202}
2203
2204// Allow private_bounds because these internal types (ToolHandlerService, TypedHandler, etc.)
2205// are implementation details that users don't interact with directly.
2206#[allow(private_bounds)]
2207impl<I, F, Fut, L> ToolBuilderWithLayer<I, F, L>
2208where
2209    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2210    F: Fn(I) -> Fut + Send + Sync + 'static,
2211    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2212    L: tower::Layer<ToolHandlerService<TypedHandler<I, F>>> + Clone + Send + Sync + 'static,
2213    L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
2214    <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
2215    <L::Service as Service<ToolRequest>>::Future: Send,
2216{
2217    /// Build the tool with the applied layer(s).
2218    pub fn build(self) -> Tool {
2219        let input_schema = self.input_schema_override.unwrap_or_else(|| {
2220            let input_schema = schemars::schema_for!(I);
2221            serde_json::to_value(input_schema)
2222                .unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
2223        });
2224        let input_schema = ensure_object_schema(input_schema);
2225
2226        let handler_service = ToolHandlerService::new(TypedHandler {
2227            handler: self.handler,
2228            _phantom: std::marker::PhantomData,
2229        });
2230        let layered = self.layer.layer(handler_service);
2231        let catch_error = ToolCatchError::new(layered);
2232        let service = BoxCloneService::new(catch_error);
2233
2234        Tool {
2235            name: self.name,
2236            title: self.title,
2237            description: self.description,
2238            output_schema: self.output_schema,
2239            icons: self.icons,
2240            annotations: self.annotations,
2241            meta: None,
2242            task_support: self.task_support,
2243            required_client_capabilities: None,
2244            task_preparer: self.task_preparer,
2245            service: Some(service),
2246            #[cfg(feature = "stateless")]
2247            mrtr_handler: None,
2248            input_schema,
2249        }
2250    }
2251
2252    /// Apply an additional Tower layer (middleware).
2253    ///
2254    /// Layers are applied in order, with earlier layers wrapping later ones.
2255    /// This means the first layer added is the outermost middleware.
2256    pub fn layer<L2>(
2257        self,
2258        layer: L2,
2259    ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<L2, L>> {
2260        ToolBuilderWithLayer {
2261            name: self.name,
2262            title: self.title,
2263            description: self.description,
2264            output_schema: self.output_schema,
2265            input_schema_override: self.input_schema_override,
2266            icons: self.icons,
2267            annotations: self.annotations,
2268            task_support: self.task_support,
2269            task_preparer: self.task_preparer,
2270            handler: self.handler,
2271            layer: tower::layer::util::Stack::new(layer, self.layer),
2272            _phantom: std::marker::PhantomData,
2273        }
2274    }
2275
2276    /// Apply a guard to this tool.
2277    ///
2278    /// See [`ToolBuilderWithHandler::guard`] for details.
2279    pub fn guard<G>(
2280        self,
2281        guard: G,
2282    ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2283    where
2284        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2285    {
2286        self.layer(GuardLayer::new(guard))
2287    }
2288}
2289
2290// =============================================================================
2291// Handler implementations
2292// =============================================================================
2293
2294/// Handler that deserializes input to a specific type
2295struct TypedHandler<I, F> {
2296    handler: F,
2297    _phantom: std::marker::PhantomData<I>,
2298}
2299
2300#[cfg(feature = "stateless")]
2301struct TypedMrtrHandler<I, F> {
2302    handler: F,
2303    _phantom: std::marker::PhantomData<I>,
2304}
2305
2306#[cfg(feature = "stateless")]
2307impl<I, F, Fut> MrtrToolHandler for TypedMrtrHandler<I, F>
2308where
2309    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2310    F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2311    Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2312{
2313    fn call(
2314        &self,
2315        ctx: RequestContext,
2316        args: Value,
2317    ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
2318        Box::pin(async move {
2319            let input: I = serde_json::from_value(args)
2320                .map_err(|error| Error::invalid_params(format!("Invalid input: {error}")))?;
2321            (self.handler)(ctx, input).await
2322        })
2323    }
2324
2325    fn input_schema(&self) -> Value {
2326        let schema = schemars::schema_for!(I);
2327        ensure_object_schema(
2328            serde_json::to_value(schema)
2329                .unwrap_or_else(|_| serde_json::json!({ "type": "object" })),
2330        )
2331    }
2332}
2333
2334impl<I, F, Fut> ToolHandler for TypedHandler<I, F>
2335where
2336    I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2337    F: Fn(I) -> Fut + Send + Sync + 'static,
2338    Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2339{
2340    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2341        Box::pin(async move {
2342            let input: I = match serde_json::from_value(args) {
2343                Ok(input) => input,
2344                Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2345            };
2346            (self.handler)(input).await
2347        })
2348    }
2349
2350    fn input_schema(&self) -> Value {
2351        let schema = schemars::schema_for!(I);
2352        let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2353            serde_json::json!({
2354                "type": "object"
2355            })
2356        });
2357        ensure_object_schema(schema)
2358    }
2359}
2360
2361// =============================================================================
2362// Trait-based tool definition
2363// =============================================================================
2364
2365/// Trait for defining tools with full control
2366///
2367/// Implement this trait when you need more control than the builder provides,
2368/// or when you want to define tools as standalone types.
2369///
2370/// # Example
2371///
2372/// ```rust
2373/// use tower_mcp::tool::McpTool;
2374/// use tower_mcp::error::Result;
2375/// use schemars::JsonSchema;
2376/// use serde::{Deserialize, Serialize};
2377///
2378/// #[derive(Debug, Deserialize, JsonSchema)]
2379/// struct AddInput {
2380///     a: i64,
2381///     b: i64,
2382/// }
2383///
2384/// struct AddTool;
2385///
2386/// impl McpTool for AddTool {
2387///     const NAME: &'static str = "add";
2388///     const DESCRIPTION: &'static str = "Add two numbers";
2389///
2390///     type Input = AddInput;
2391///     type Output = i64;
2392///
2393///     async fn call(&self, input: Self::Input) -> Result<Self::Output> {
2394///         Ok(input.a + input.b)
2395///     }
2396/// }
2397///
2398/// let tool = AddTool.into_tool();
2399/// assert_eq!(tool.name, "add");
2400/// ```
2401pub trait McpTool: Send + Sync + 'static {
2402    /// The tool name (must be unique within the router).
2403    const NAME: &'static str;
2404    /// A human-readable description of the tool.
2405    const DESCRIPTION: &'static str;
2406
2407    /// The input type, deserialized from tool call arguments.
2408    type Input: JsonSchema + DeserializeOwned + Send;
2409    /// The output type, serialized into the tool call result.
2410    type Output: Serialize + Send;
2411
2412    /// Execute the tool with the given input.
2413    fn call(&self, input: Self::Input) -> impl Future<Output = Result<Self::Output>> + Send;
2414
2415    /// Optional annotations for the tool
2416    fn annotations(&self) -> Option<ToolAnnotations> {
2417        None
2418    }
2419
2420    /// Convert to a [`Tool`] instance.
2421    ///
2422    /// # Panics
2423    ///
2424    /// Panics if [`NAME`](Self::NAME) is not a valid tool name. Since `NAME`
2425    /// is a `&'static str`, invalid names are caught immediately during
2426    /// development.
2427    fn into_tool(self) -> Tool
2428    where
2429        Self: Sized,
2430    {
2431        if let Err(e) = validate_tool_name(Self::NAME) {
2432            panic!("{e}");
2433        }
2434        let annotations = self.annotations();
2435        let tool = Arc::new(self);
2436        Tool::from_handler(
2437            Self::NAME.to_string(),
2438            None,
2439            Some(Self::DESCRIPTION.to_string()),
2440            None,
2441            None,
2442            annotations,
2443            TaskSupportMode::default(),
2444            None,
2445            McpToolHandler { tool },
2446        )
2447    }
2448}
2449
2450/// Wrapper to make McpTool implement ToolHandler
2451struct McpToolHandler<T: McpTool> {
2452    tool: Arc<T>,
2453}
2454
2455impl<T: McpTool> ToolHandler for McpToolHandler<T> {
2456    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2457        let tool = self.tool.clone();
2458        Box::pin(async move {
2459            let input: T::Input = match serde_json::from_value(args) {
2460                Ok(input) => input,
2461                Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2462            };
2463            let output = tool.call(input).await?;
2464            let value = serde_json::to_value(output).tool_context("Failed to serialize output")?;
2465            Ok(CallToolResult::json(value))
2466        })
2467    }
2468
2469    fn input_schema(&self) -> Value {
2470        let schema = schemars::schema_for!(T::Input);
2471        let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2472            serde_json::json!({
2473                "type": "object"
2474            })
2475        });
2476        ensure_object_schema(schema)
2477    }
2478}
2479
2480#[cfg(test)]
2481mod tests {
2482    use super::*;
2483    use crate::extract::{Context, Json, RawArgs, State};
2484    use crate::protocol::Content;
2485    use schemars::JsonSchema;
2486    use serde::Deserialize;
2487
2488    #[derive(Debug, Deserialize, JsonSchema)]
2489    struct GreetInput {
2490        name: String,
2491    }
2492
2493    #[tokio::test]
2494    async fn test_builder_tool() {
2495        let tool = ToolBuilder::new("greet")
2496            .description("Greet someone")
2497            .handler(|input: GreetInput| async move {
2498                Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2499            })
2500            .build();
2501
2502        assert_eq!(tool.name, "greet");
2503        assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2504
2505        let result = tool.call(serde_json::json!({"name": "World"})).await;
2506
2507        assert!(!result.is_error);
2508    }
2509
2510    #[cfg(feature = "stateless")]
2511    #[tokio::test]
2512    async fn test_mrtr_builder_preserves_input_required_outcome() {
2513        let tool = ToolBuilder::new("continue")
2514            .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2515                Ok(RequestOutcome::input_required(
2516                    crate::protocol::InputRequiredResult::new().with_request_state("signed-state"),
2517                ))
2518            })
2519            .build();
2520
2521        let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2522        assert_eq!(
2523            outcome
2524                .as_input_required()
2525                .and_then(|result| result.request_state.as_deref()),
2526            Some("signed-state")
2527        );
2528    }
2529
2530    #[cfg(feature = "stateless")]
2531    #[tokio::test]
2532    async fn mrtr_builder_composes_guards_and_layers() {
2533        use std::sync::atomic::{AtomicUsize, Ordering};
2534        use std::time::Duration;
2535        use tower::timeout::TimeoutLayer;
2536
2537        let rounds = Arc::new(AtomicUsize::new(0));
2538        let observed = rounds.clone();
2539        let tool = ToolBuilder::new("guarded_continue")
2540            .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2541                Ok(RequestOutcome::input_required(
2542                    crate::protocol::InputRequiredResult::new().with_request_state("continue"),
2543                ))
2544            })
2545            .layer(TimeoutLayer::new(Duration::from_secs(1)))
2546            .guard(move |_request| {
2547                observed.fetch_add(1, Ordering::SeqCst);
2548                Ok(())
2549            })
2550            .build();
2551
2552        for _ in 0..2 {
2553            assert!(
2554                tool.call_outcome(serde_json::json!({}))
2555                    .await
2556                    .unwrap()
2557                    .as_input_required()
2558                    .is_some()
2559            );
2560        }
2561        assert_eq!(rounds.load(Ordering::SeqCst), 2);
2562    }
2563
2564    #[cfg(feature = "stateless")]
2565    #[tokio::test]
2566    async fn built_mrtr_tool_accepts_a_guard() {
2567        let tool = ToolBuilder::new("denied_continue")
2568            .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2569                Ok(RequestOutcome::input_required(
2570                    crate::protocol::InputRequiredResult::new().with_request_state("unreachable"),
2571                ))
2572            })
2573            .build()
2574            .with_guard(|_request| Err("MRTR access denied".to_string()));
2575
2576        let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2577        let result = outcome
2578            .as_complete()
2579            .expect("guard rejection is a complete tool error");
2580        assert!(result.is_error);
2581        assert_eq!(result.first_text(), Some("MRTR access denied"));
2582    }
2583
2584    #[tokio::test]
2585    async fn test_raw_handler() {
2586        let tool = ToolBuilder::new("echo")
2587            .description("Echo input")
2588            .extractor_handler((), |RawArgs(args): RawArgs| async move {
2589                Ok(CallToolResult::json(args))
2590            })
2591            .build();
2592
2593        let result = tool.call(serde_json::json!({"foo": "bar"})).await;
2594
2595        assert!(!result.is_error);
2596    }
2597
2598    #[test]
2599    fn test_invalid_tool_name_empty() {
2600        let err = ToolBuilder::try_new("").err().expect("should fail");
2601        assert!(err.to_string().contains("cannot be empty"));
2602    }
2603
2604    #[test]
2605    fn test_invalid_tool_name_too_long() {
2606        let long_name = "a".repeat(65);
2607        let err = ToolBuilder::try_new(long_name).err().expect("should fail");
2608        assert!(err.to_string().contains("exceeds maximum"));
2609    }
2610
2611    #[test]
2612    fn test_invalid_tool_name_bad_chars() {
2613        let err = ToolBuilder::try_new("my tool!").err().expect("should fail");
2614        assert!(err.to_string().contains("invalid character"));
2615    }
2616
2617    #[test]
2618    #[should_panic(expected = "cannot be empty")]
2619    fn test_new_panics_on_empty_name() {
2620        ToolBuilder::new("");
2621    }
2622
2623    #[test]
2624    #[should_panic(expected = "exceeds maximum")]
2625    fn test_new_panics_on_too_long_name() {
2626        ToolBuilder::new("a".repeat(65));
2627    }
2628
2629    #[test]
2630    #[should_panic(expected = "invalid character")]
2631    fn test_new_panics_on_invalid_chars() {
2632        ToolBuilder::new("my tool!");
2633    }
2634
2635    #[test]
2636    fn test_valid_tool_names() {
2637        // All valid characters per SEP-986
2638        let names = [
2639            "my_tool",
2640            "my-tool",
2641            "my.tool",
2642            "my/tool",
2643            "user-profile/update",
2644            "MyTool123",
2645            "a",
2646            &"a".repeat(64),
2647        ];
2648        for name in names {
2649            assert!(
2650                ToolBuilder::try_new(name).is_ok(),
2651                "Expected '{}' to be valid",
2652                name
2653            );
2654        }
2655    }
2656
2657    #[tokio::test]
2658    async fn test_context_aware_handler() {
2659        use crate::context::notification_channel;
2660        use crate::protocol::{ProgressToken, RequestId};
2661
2662        #[derive(Debug, Deserialize, JsonSchema)]
2663        struct ProcessInput {
2664            count: i32,
2665        }
2666
2667        let tool = ToolBuilder::new("process")
2668            .description("Process with context")
2669            .extractor_handler(
2670                (),
2671                |ctx: Context, Json(input): Json<ProcessInput>| async move {
2672                    // Simulate progress reporting
2673                    for i in 0..input.count {
2674                        if ctx.is_cancelled() {
2675                            return Ok(CallToolResult::error("Cancelled"));
2676                        }
2677                        ctx.report_progress(i as f64, Some(input.count as f64), None)
2678                            .await;
2679                    }
2680                    Ok(CallToolResult::text(format!(
2681                        "Processed {} items",
2682                        input.count
2683                    )))
2684                },
2685            )
2686            .build();
2687
2688        assert_eq!(tool.name, "process");
2689
2690        // Test with a context that has progress token and notification sender
2691        let (tx, mut rx) = notification_channel(10);
2692        let ctx = RequestContext::new(RequestId::Number(1))
2693            .with_progress_token(ProgressToken::Number(42))
2694            .with_notification_sender(tx);
2695
2696        let result = tool
2697            .call_with_context(ctx, serde_json::json!({"count": 3}))
2698            .await;
2699
2700        assert!(!result.is_error);
2701
2702        // Check that progress notifications were sent
2703        let mut progress_count = 0;
2704        while rx.try_recv().is_ok() {
2705            progress_count += 1;
2706        }
2707        assert_eq!(progress_count, 3);
2708    }
2709
2710    #[tokio::test]
2711    async fn test_context_aware_handler_cancellation() {
2712        use crate::protocol::RequestId;
2713        use std::sync::atomic::{AtomicI32, Ordering};
2714
2715        #[derive(Debug, Deserialize, JsonSchema)]
2716        struct LongRunningInput {
2717            iterations: i32,
2718        }
2719
2720        let iterations_completed = Arc::new(AtomicI32::new(0));
2721        let iterations_ref = iterations_completed.clone();
2722
2723        let tool = ToolBuilder::new("long_running")
2724            .description("Long running task")
2725            .extractor_handler(
2726                (),
2727                move |ctx: Context, Json(input): Json<LongRunningInput>| {
2728                    let completed = iterations_ref.clone();
2729                    async move {
2730                        for i in 0..input.iterations {
2731                            if ctx.is_cancelled() {
2732                                return Ok(CallToolResult::error("Cancelled"));
2733                            }
2734                            completed.fetch_add(1, Ordering::SeqCst);
2735                            // Simulate work
2736                            tokio::task::yield_now().await;
2737                            // Cancel after iteration 2
2738                            if i == 2 {
2739                                ctx.cancellation_token().cancel();
2740                            }
2741                        }
2742                        Ok(CallToolResult::text("Done"))
2743                    }
2744                },
2745            )
2746            .build();
2747
2748        let ctx = RequestContext::new(RequestId::Number(1));
2749
2750        let result = tool
2751            .call_with_context(ctx, serde_json::json!({"iterations": 10}))
2752            .await;
2753
2754        // Should have been cancelled after 3 iterations (0, 1, 2)
2755        // The next iteration (3) checks cancellation and returns
2756        assert!(result.is_error);
2757        assert_eq!(iterations_completed.load(Ordering::SeqCst), 3);
2758    }
2759
2760    #[tokio::test]
2761    async fn test_tool_builder_with_enhanced_fields() {
2762        let output_schema = serde_json::json!({
2763            "type": "object",
2764            "properties": {
2765                "greeting": {"type": "string"}
2766            }
2767        });
2768
2769        let tool = ToolBuilder::new("greet")
2770            .title("Greeting Tool")
2771            .description("Greet someone")
2772            .output_schema(output_schema.clone())
2773            .icon("https://example.com/icon.png")
2774            .icon_with_meta(
2775                "https://example.com/icon-large.png",
2776                Some("image/png".to_string()),
2777                Some(vec!["96x96".to_string()]),
2778            )
2779            .handler(|input: GreetInput| async move {
2780                Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2781            })
2782            .build();
2783
2784        assert_eq!(tool.name, "greet");
2785        assert_eq!(tool.title.as_deref(), Some("Greeting Tool"));
2786        assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2787        assert_eq!(tool.output_schema, Some(output_schema));
2788        assert!(tool.icons.is_some());
2789        assert_eq!(tool.icons.as_ref().unwrap().len(), 2);
2790
2791        // Test definition includes new fields
2792        let def = tool.definition();
2793        assert_eq!(def.title.as_deref(), Some("Greeting Tool"));
2794        assert!(def.output_schema.is_some());
2795        assert!(def.icons.is_some());
2796    }
2797
2798    #[tokio::test]
2799    async fn test_handler_with_state() {
2800        let shared = Arc::new("shared-state".to_string());
2801
2802        let tool = ToolBuilder::new("stateful")
2803            .description("Uses shared state")
2804            .extractor_handler(
2805                shared,
2806                |State(state): State<Arc<String>>, Json(input): Json<GreetInput>| async move {
2807                    Ok(CallToolResult::text(format!(
2808                        "{}: Hello, {}!",
2809                        state, input.name
2810                    )))
2811                },
2812            )
2813            .build();
2814
2815        let result = tool.call(serde_json::json!({"name": "World"})).await;
2816        assert!(!result.is_error);
2817    }
2818
2819    #[tokio::test]
2820    async fn test_handler_with_state_and_context() {
2821        use crate::protocol::RequestId;
2822
2823        let shared = Arc::new(42_i32);
2824
2825        let tool =
2826            ToolBuilder::new("stateful_ctx")
2827                .description("Uses state and context")
2828                .extractor_handler(
2829                    shared,
2830                    |State(state): State<Arc<i32>>,
2831                     _ctx: Context,
2832                     Json(input): Json<GreetInput>| async move {
2833                        Ok(CallToolResult::text(format!(
2834                            "{}: Hello, {}!",
2835                            state, input.name
2836                        )))
2837                    },
2838                )
2839                .build();
2840
2841        let ctx = RequestContext::new(RequestId::Number(1));
2842        let result = tool
2843            .call_with_context(ctx, serde_json::json!({"name": "World"}))
2844            .await;
2845        assert!(!result.is_error);
2846    }
2847
2848    #[tokio::test]
2849    async fn test_handler_no_params() {
2850        let tool = ToolBuilder::new("no_params")
2851            .description("Takes no parameters")
2852            .extractor_handler((), |Json(_): Json<NoParams>| async {
2853                Ok(CallToolResult::text("no params result"))
2854            })
2855            .build();
2856
2857        assert_eq!(tool.name, "no_params");
2858
2859        // Should work with empty args
2860        let result = tool.call(serde_json::json!({})).await;
2861        assert!(!result.is_error);
2862
2863        // Should also work with unexpected args (ignored)
2864        let result = tool.call(serde_json::json!({"unexpected": "value"})).await;
2865        assert!(!result.is_error);
2866
2867        // Check input schema includes type: object
2868        let schema = tool.definition().input_schema;
2869        assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2870    }
2871
2872    #[tokio::test]
2873    async fn test_handler_with_state_no_params() {
2874        let shared = Arc::new("shared_value".to_string());
2875
2876        let tool = ToolBuilder::new("with_state_no_params")
2877            .description("Takes no parameters but has state")
2878            .extractor_handler(
2879                shared,
2880                |State(state): State<Arc<String>>, Json(_): Json<NoParams>| async move {
2881                    Ok(CallToolResult::text(format!("state: {}", state)))
2882                },
2883            )
2884            .build();
2885
2886        assert_eq!(tool.name, "with_state_no_params");
2887
2888        // Should work with empty args
2889        let result = tool.call(serde_json::json!({})).await;
2890        assert!(!result.is_error);
2891        assert_eq!(result.first_text().unwrap(), "state: shared_value");
2892
2893        // Check input schema includes type: object
2894        let schema = tool.definition().input_schema;
2895        assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2896    }
2897
2898    #[tokio::test]
2899    async fn test_handler_no_params_with_context() {
2900        let tool = ToolBuilder::new("no_params_with_context")
2901            .description("Takes no parameters but has context")
2902            .extractor_handler((), |_ctx: Context, Json(_): Json<NoParams>| async move {
2903                Ok(CallToolResult::text("context available"))
2904            })
2905            .build();
2906
2907        assert_eq!(tool.name, "no_params_with_context");
2908
2909        let result = tool.call(serde_json::json!({})).await;
2910        assert!(!result.is_error);
2911        assert_eq!(result.first_text().unwrap(), "context available");
2912    }
2913
2914    #[tokio::test]
2915    async fn test_handler_with_state_and_context_no_params() {
2916        let shared = Arc::new("shared".to_string());
2917
2918        let tool = ToolBuilder::new("state_context_no_params")
2919            .description("Has state and context, no params")
2920            .extractor_handler(
2921                shared,
2922                |State(state): State<Arc<String>>,
2923                 _ctx: Context,
2924                 Json(_): Json<NoParams>| async move {
2925                    Ok(CallToolResult::text(format!("state: {}", state)))
2926                },
2927            )
2928            .build();
2929
2930        assert_eq!(tool.name, "state_context_no_params");
2931
2932        let result = tool.call(serde_json::json!({})).await;
2933        assert!(!result.is_error);
2934        assert_eq!(result.first_text().unwrap(), "state: shared");
2935    }
2936
2937    #[tokio::test]
2938    async fn test_raw_handler_with_state() {
2939        let prefix = Arc::new("prefix:".to_string());
2940
2941        let tool = ToolBuilder::new("raw_with_state")
2942            .description("Raw handler with state")
2943            .extractor_handler(
2944                prefix,
2945                |State(state): State<Arc<String>>, RawArgs(args): RawArgs| async move {
2946                    Ok(CallToolResult::text(format!("{} {}", state, args)))
2947                },
2948            )
2949            .build();
2950
2951        assert_eq!(tool.name, "raw_with_state");
2952
2953        let result = tool.call(serde_json::json!({"key": "value"})).await;
2954        assert!(!result.is_error);
2955        assert!(result.first_text().unwrap().starts_with("prefix:"));
2956    }
2957
2958    #[tokio::test]
2959    async fn test_raw_handler_with_state_and_context() {
2960        let prefix = Arc::new("prefix:".to_string());
2961
2962        let tool = ToolBuilder::new("raw_state_context")
2963            .description("Raw handler with state and context")
2964            .extractor_handler(
2965                prefix,
2966                |State(state): State<Arc<String>>,
2967                 _ctx: Context,
2968                 RawArgs(args): RawArgs| async move {
2969                    Ok(CallToolResult::text(format!("{} {}", state, args)))
2970                },
2971            )
2972            .build();
2973
2974        assert_eq!(tool.name, "raw_state_context");
2975
2976        let result = tool.call(serde_json::json!({"key": "value"})).await;
2977        assert!(!result.is_error);
2978        assert!(result.first_text().unwrap().starts_with("prefix:"));
2979    }
2980
2981    #[tokio::test]
2982    async fn test_tool_with_timeout_layer() {
2983        use std::time::Duration;
2984        use tower::timeout::TimeoutLayer;
2985
2986        #[derive(Debug, Deserialize, JsonSchema)]
2987        struct SlowInput {
2988            delay_ms: u64,
2989        }
2990
2991        // Create a tool with a short timeout
2992        let tool = ToolBuilder::new("slow_tool")
2993            .description("A slow tool")
2994            .handler(|input: SlowInput| async move {
2995                tokio::time::sleep(Duration::from_millis(input.delay_ms)).await;
2996                Ok(CallToolResult::text("completed"))
2997            })
2998            .layer(TimeoutLayer::new(Duration::from_millis(50)))
2999            .build();
3000
3001        // Fast call should succeed
3002        let result = tool.call(serde_json::json!({"delay_ms": 10})).await;
3003        assert!(!result.is_error);
3004        assert_eq!(result.first_text().unwrap(), "completed");
3005
3006        // Slow call should timeout and return an error result
3007        let result = tool.call(serde_json::json!({"delay_ms": 200})).await;
3008        assert!(result.is_error);
3009        // Tower's timeout error message is "request timed out"
3010        let msg = result.first_text().unwrap().to_lowercase();
3011        assert!(
3012            msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
3013            "Expected timeout error, got: {}",
3014            msg
3015        );
3016    }
3017
3018    #[tokio::test]
3019    async fn test_tool_with_concurrency_limit_layer() {
3020        use std::sync::atomic::{AtomicU32, Ordering};
3021        use std::time::Duration;
3022        use tower::limit::ConcurrencyLimitLayer;
3023
3024        #[derive(Debug, Deserialize, JsonSchema)]
3025        struct WorkInput {
3026            id: u32,
3027        }
3028
3029        let max_concurrent = Arc::new(AtomicU32::new(0));
3030        let current_concurrent = Arc::new(AtomicU32::new(0));
3031        let max_ref = max_concurrent.clone();
3032        let current_ref = current_concurrent.clone();
3033
3034        // Create a tool with concurrency limit of 2
3035        let tool = ToolBuilder::new("concurrent_tool")
3036            .description("A concurrent tool")
3037            .handler(move |input: WorkInput| {
3038                let max = max_ref.clone();
3039                let current = current_ref.clone();
3040                async move {
3041                    // Track concurrency
3042                    let prev = current.fetch_add(1, Ordering::SeqCst);
3043                    max.fetch_max(prev + 1, Ordering::SeqCst);
3044
3045                    // Simulate work
3046                    tokio::time::sleep(Duration::from_millis(50)).await;
3047
3048                    current.fetch_sub(1, Ordering::SeqCst);
3049                    Ok(CallToolResult::text(format!("completed {}", input.id)))
3050                }
3051            })
3052            .layer(ConcurrencyLimitLayer::new(2))
3053            .build();
3054
3055        // Launch 4 concurrent calls
3056        let handles: Vec<_> = (0..4)
3057            .map(|i| {
3058                let t = tool.call(serde_json::json!({"id": i}));
3059                tokio::spawn(t)
3060            })
3061            .collect();
3062
3063        for handle in handles {
3064            let result = handle.await.unwrap();
3065            assert!(!result.is_error);
3066        }
3067
3068        // Max concurrent should not exceed 2
3069        assert!(max_concurrent.load(Ordering::SeqCst) <= 2);
3070    }
3071
3072    #[tokio::test]
3073    async fn test_tool_with_multiple_layers() {
3074        use std::time::Duration;
3075        use tower::limit::ConcurrencyLimitLayer;
3076        use tower::timeout::TimeoutLayer;
3077
3078        #[derive(Debug, Deserialize, JsonSchema)]
3079        struct Input {
3080            value: String,
3081        }
3082
3083        // Create a tool with multiple layers stacked
3084        let tool = ToolBuilder::new("multi_layer_tool")
3085            .description("Tool with multiple layers")
3086            .handler(|input: Input| async move {
3087                Ok(CallToolResult::text(format!("processed: {}", input.value)))
3088            })
3089            .layer(TimeoutLayer::new(Duration::from_secs(5)))
3090            .layer(ConcurrencyLimitLayer::new(10))
3091            .build();
3092
3093        let result = tool.call(serde_json::json!({"value": "test"})).await;
3094        assert!(!result.is_error);
3095        assert_eq!(result.first_text().unwrap(), "processed: test");
3096    }
3097
3098    #[test]
3099    fn test_tool_catch_error_clone() {
3100        // ToolCatchError should be Clone when inner is Clone
3101        // Use a simple tool that we can clone
3102        let tool = ToolBuilder::new("test")
3103            .description("test")
3104            .extractor_handler((), |RawArgs(_args): RawArgs| async {
3105                Ok(CallToolResult::text("ok"))
3106            })
3107            .build();
3108        // The tool contains a BoxToolService which is cloneable
3109        let _clone = tool.call(serde_json::json!({}));
3110    }
3111
3112    #[test]
3113    fn test_tool_catch_error_debug() {
3114        // ToolCatchError implements Debug when inner implements Debug
3115        // Since our internal services don't require Debug, just verify
3116        // that ToolCatchError has a Debug impl for appropriate types
3117        #[derive(Debug, Clone)]
3118        struct DebugService;
3119
3120        impl Service<ToolRequest> for DebugService {
3121            type Response = CallToolResult;
3122            type Error = crate::error::Error;
3123            type Future = Pin<
3124                Box<
3125                    dyn Future<Output = std::result::Result<CallToolResult, crate::error::Error>>
3126                        + Send,
3127                >,
3128            >;
3129
3130            fn poll_ready(
3131                &mut self,
3132                _cx: &mut std::task::Context<'_>,
3133            ) -> Poll<std::result::Result<(), Self::Error>> {
3134                Poll::Ready(Ok(()))
3135            }
3136
3137            fn call(&mut self, _req: ToolRequest) -> Self::Future {
3138                Box::pin(async { Ok(CallToolResult::text("ok")) })
3139            }
3140        }
3141
3142        let catch_error = ToolCatchError::new(DebugService);
3143        let debug = format!("{:?}", catch_error);
3144        assert!(debug.contains("ToolCatchError"));
3145    }
3146
3147    #[test]
3148    fn test_tool_request_new() {
3149        use crate::protocol::RequestId;
3150
3151        let ctx = RequestContext::new(RequestId::Number(42));
3152        let args = serde_json::json!({"key": "value"});
3153        let req = ToolRequest::new(ctx.clone(), args.clone());
3154
3155        assert_eq!(req.args, args);
3156    }
3157
3158    #[test]
3159    fn test_no_params_schema() {
3160        // NoParams should produce a schema with type: "object"
3161        let schema = schemars::schema_for!(NoParams);
3162        let schema_value = serde_json::to_value(&schema).unwrap();
3163        assert_eq!(
3164            schema_value.get("type").and_then(|v| v.as_str()),
3165            Some("object"),
3166            "NoParams should generate type: object schema"
3167        );
3168    }
3169
3170    #[test]
3171    fn test_no_params_deserialize() {
3172        // NoParams should deserialize from various inputs
3173        let from_empty_object: NoParams = serde_json::from_str("{}").unwrap();
3174        assert_eq!(from_empty_object, NoParams);
3175
3176        let from_null: NoParams = serde_json::from_str("null").unwrap();
3177        assert_eq!(from_null, NoParams);
3178
3179        // Should also accept objects with unexpected fields (ignored)
3180        let from_object_with_fields: NoParams =
3181            serde_json::from_str(r#"{"unexpected": "value"}"#).unwrap();
3182        assert_eq!(from_object_with_fields, NoParams);
3183    }
3184
3185    #[tokio::test]
3186    async fn test_no_params_type_in_handler() {
3187        // NoParams can be used as a handler input type
3188        let tool = ToolBuilder::new("status")
3189            .description("Get status")
3190            .handler(|_input: NoParams| async move { Ok(CallToolResult::text("OK")) })
3191            .build();
3192
3193        // Check schema has type: object (not type: null like () would produce)
3194        let schema = tool.definition().input_schema;
3195        assert_eq!(
3196            schema.get("type").and_then(|v| v.as_str()),
3197            Some("object"),
3198            "NoParams handler should produce type: object schema"
3199        );
3200
3201        // Should work with empty input
3202        let result = tool.call(serde_json::json!({})).await;
3203        assert!(!result.is_error);
3204    }
3205
3206    #[tokio::test]
3207    async fn test_serde_json_value_handler_has_type_object() {
3208        // serde_json::Value generates a schema without "type" via schemars.
3209        // We must ensure "type": "object" is added for MCP compliance.
3210        let tool = ToolBuilder::new("any_input")
3211            .description("Accepts any input")
3212            .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3213            .build();
3214
3215        let schema = tool.definition().input_schema;
3216        assert_eq!(
3217            schema.get("type").and_then(|v| v.as_str()),
3218            Some("object"),
3219            "serde_json::Value handler should produce schema with type: object"
3220        );
3221    }
3222
3223    #[tokio::test]
3224    async fn test_tool_with_name_prefix() {
3225        #[derive(Debug, Deserialize, JsonSchema)]
3226        struct Input {
3227            value: String,
3228        }
3229
3230        let tool = ToolBuilder::new("query")
3231            .description("Query something")
3232            .title("Query Tool")
3233            .handler(|input: Input| async move { Ok(CallToolResult::text(&input.value)) })
3234            .build();
3235
3236        // Create prefixed version
3237        let prefixed = tool.with_name_prefix("db");
3238
3239        // Check name is prefixed
3240        assert_eq!(prefixed.name, "db.query");
3241
3242        // Check other fields are preserved
3243        assert_eq!(prefixed.description.as_deref(), Some("Query something"));
3244        assert_eq!(prefixed.title.as_deref(), Some("Query Tool"));
3245
3246        // Check the tool still works
3247        let result = prefixed
3248            .call(serde_json::json!({"value": "test input"}))
3249            .await;
3250        assert!(!result.is_error);
3251        match &result.content[0] {
3252            Content::Text { text, .. } => assert_eq!(text, "test input"),
3253            _ => panic!("Expected text content"),
3254        }
3255    }
3256
3257    #[tokio::test]
3258    async fn test_tool_with_name_prefix_multiple_levels() {
3259        let tool = ToolBuilder::new("action")
3260            .description("Do something")
3261            .handler(|_: NoParams| async move { Ok(CallToolResult::text("done")) })
3262            .build();
3263
3264        // Apply multiple prefixes
3265        let prefixed = tool.with_name_prefix("level1");
3266        assert_eq!(prefixed.name, "level1.action");
3267
3268        let double_prefixed = prefixed.with_name_prefix("level0");
3269        assert_eq!(double_prefixed.name, "level0.level1.action");
3270    }
3271
3272    // =============================================================================
3273    // no_params_handler tests
3274    // =============================================================================
3275
3276    #[tokio::test]
3277    async fn test_no_params_handler_basic() {
3278        let tool = ToolBuilder::new("get_status")
3279            .description("Get current status")
3280            .no_params_handler(|| async { Ok(CallToolResult::text("OK")) })
3281            .build();
3282
3283        assert_eq!(tool.name, "get_status");
3284        assert_eq!(tool.description.as_deref(), Some("Get current status"));
3285
3286        // Should work with empty args
3287        let result = tool.call(serde_json::json!({})).await;
3288        assert!(!result.is_error);
3289        assert_eq!(result.first_text().unwrap(), "OK");
3290
3291        // Should also work with null args
3292        let result = tool.call(serde_json::json!(null)).await;
3293        assert!(!result.is_error);
3294
3295        // Check input schema has type: object
3296        let schema = tool.definition().input_schema;
3297        assert_eq!(schema.get("type").and_then(|v| v.as_str()), Some("object"));
3298    }
3299
3300    #[tokio::test]
3301    async fn test_no_params_handler_with_captured_state() {
3302        let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
3303        let counter_ref = counter.clone();
3304
3305        let tool = ToolBuilder::new("increment")
3306            .description("Increment counter")
3307            .no_params_handler(move || {
3308                let c = counter_ref.clone();
3309                async move {
3310                    let prev = c.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
3311                    Ok(CallToolResult::text(format!("Incremented from {}", prev)))
3312                }
3313            })
3314            .build();
3315
3316        // Call multiple times
3317        let _ = tool.call(serde_json::json!({})).await;
3318        let _ = tool.call(serde_json::json!({})).await;
3319        let result = tool.call(serde_json::json!({})).await;
3320
3321        assert!(!result.is_error);
3322        assert_eq!(result.first_text().unwrap(), "Incremented from 2");
3323        assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 3);
3324    }
3325
3326    #[tokio::test]
3327    async fn test_no_params_handler_with_layer() {
3328        use std::time::Duration;
3329        use tower::timeout::TimeoutLayer;
3330
3331        let tool = ToolBuilder::new("slow_status")
3332            .description("Slow status check")
3333            .no_params_handler(|| async {
3334                tokio::time::sleep(Duration::from_millis(10)).await;
3335                Ok(CallToolResult::text("done"))
3336            })
3337            .layer(TimeoutLayer::new(Duration::from_secs(1)))
3338            .build();
3339
3340        let result = tool.call(serde_json::json!({})).await;
3341        assert!(!result.is_error);
3342        assert_eq!(result.first_text().unwrap(), "done");
3343    }
3344
3345    #[tokio::test]
3346    async fn test_no_params_handler_timeout() {
3347        use std::time::Duration;
3348        use tower::timeout::TimeoutLayer;
3349
3350        let tool = ToolBuilder::new("very_slow_status")
3351            .description("Very slow status check")
3352            .no_params_handler(|| async {
3353                tokio::time::sleep(Duration::from_millis(200)).await;
3354                Ok(CallToolResult::text("done"))
3355            })
3356            .layer(TimeoutLayer::new(Duration::from_millis(50)))
3357            .build();
3358
3359        let result = tool.call(serde_json::json!({})).await;
3360        assert!(result.is_error);
3361        let msg = result.first_text().unwrap().to_lowercase();
3362        assert!(
3363            msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
3364            "Expected timeout error, got: {}",
3365            msg
3366        );
3367    }
3368
3369    #[tokio::test]
3370    async fn test_no_params_handler_with_multiple_layers() {
3371        use std::time::Duration;
3372        use tower::limit::ConcurrencyLimitLayer;
3373        use tower::timeout::TimeoutLayer;
3374
3375        let tool = ToolBuilder::new("multi_layer_status")
3376            .description("Status with multiple layers")
3377            .no_params_handler(|| async { Ok(CallToolResult::text("status ok")) })
3378            .layer(TimeoutLayer::new(Duration::from_secs(5)))
3379            .layer(ConcurrencyLimitLayer::new(10))
3380            .build();
3381
3382        let result = tool.call(serde_json::json!({})).await;
3383        assert!(!result.is_error);
3384        assert_eq!(result.first_text().unwrap(), "status ok");
3385    }
3386
3387    // =========================================================================
3388    // Guard tests
3389    // =========================================================================
3390
3391    #[tokio::test]
3392    async fn test_guard_allows_request() {
3393        #[derive(Debug, Deserialize, JsonSchema)]
3394        #[allow(dead_code)]
3395        struct DeleteInput {
3396            id: String,
3397            confirm: bool,
3398        }
3399
3400        let tool = ToolBuilder::new("delete")
3401            .description("Delete a record")
3402            .handler(|input: DeleteInput| async move {
3403                Ok(CallToolResult::text(format!("deleted {}", input.id)))
3404            })
3405            .guard(|req: &ToolRequest| {
3406                let confirm = req
3407                    .args
3408                    .get("confirm")
3409                    .and_then(|v| v.as_bool())
3410                    .unwrap_or(false);
3411                if !confirm {
3412                    return Err("Must set confirm=true to delete".to_string());
3413                }
3414                Ok(())
3415            })
3416            .build();
3417
3418        let result = tool
3419            .call(serde_json::json!({"id": "abc", "confirm": true}))
3420            .await;
3421        assert!(!result.is_error);
3422        assert_eq!(result.first_text().unwrap(), "deleted abc");
3423    }
3424
3425    #[tokio::test]
3426    async fn test_guard_rejects_request() {
3427        #[derive(Debug, Deserialize, JsonSchema)]
3428        #[allow(dead_code)]
3429        struct DeleteInput2 {
3430            id: String,
3431            confirm: bool,
3432        }
3433
3434        let tool = ToolBuilder::new("delete2")
3435            .description("Delete a record")
3436            .handler(|input: DeleteInput2| async move {
3437                Ok(CallToolResult::text(format!("deleted {}", input.id)))
3438            })
3439            .guard(|req: &ToolRequest| {
3440                let confirm = req
3441                    .args
3442                    .get("confirm")
3443                    .and_then(|v| v.as_bool())
3444                    .unwrap_or(false);
3445                if !confirm {
3446                    return Err("Must set confirm=true to delete".to_string());
3447                }
3448                Ok(())
3449            })
3450            .build();
3451
3452        let result = tool
3453            .call(serde_json::json!({"id": "abc", "confirm": false}))
3454            .await;
3455        assert!(result.is_error);
3456        assert!(
3457            result
3458                .first_text()
3459                .unwrap()
3460                .contains("Must set confirm=true")
3461        );
3462    }
3463
3464    #[tokio::test]
3465    async fn test_guard_with_layer() {
3466        use std::time::Duration;
3467        use tower::timeout::TimeoutLayer;
3468
3469        let tool = ToolBuilder::new("guarded_timeout")
3470            .description("Guarded with timeout")
3471            .handler(|input: GreetInput| async move {
3472                Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3473            })
3474            .layer(TimeoutLayer::new(Duration::from_secs(5)))
3475            .guard(|_req: &ToolRequest| Ok(()))
3476            .build();
3477
3478        let result = tool.call(serde_json::json!({"name": "World"})).await;
3479        assert!(!result.is_error);
3480        assert_eq!(result.first_text().unwrap(), "Hello, World!");
3481    }
3482
3483    #[tokio::test]
3484    async fn test_guard_on_no_params_handler() {
3485        let allowed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
3486        let allowed_clone = allowed.clone();
3487
3488        let tool = ToolBuilder::new("status")
3489            .description("Get status")
3490            .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3491            .guard(move |_req: &ToolRequest| {
3492                if allowed_clone.load(std::sync::atomic::Ordering::Relaxed) {
3493                    Ok(())
3494                } else {
3495                    Err("Access denied".to_string())
3496                }
3497            })
3498            .build();
3499
3500        // Allowed
3501        let result = tool.call(serde_json::json!({})).await;
3502        assert!(!result.is_error);
3503        assert_eq!(result.first_text().unwrap(), "ok");
3504
3505        // Denied
3506        allowed.store(false, std::sync::atomic::Ordering::Relaxed);
3507        let result = tool.call(serde_json::json!({})).await;
3508        assert!(result.is_error);
3509        assert!(result.first_text().unwrap().contains("Access denied"));
3510    }
3511
3512    #[tokio::test]
3513    async fn test_guard_on_no_params_handler_with_layer() {
3514        use std::time::Duration;
3515        use tower::timeout::TimeoutLayer;
3516
3517        let tool = ToolBuilder::new("status_layered")
3518            .description("Get status with layers")
3519            .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3520            .layer(TimeoutLayer::new(Duration::from_secs(5)))
3521            .guard(|_req: &ToolRequest| Ok(()))
3522            .build();
3523
3524        let result = tool.call(serde_json::json!({})).await;
3525        assert!(!result.is_error);
3526        assert_eq!(result.first_text().unwrap(), "ok");
3527    }
3528
3529    #[tokio::test]
3530    async fn test_guard_on_extractor_handler() {
3531        use std::sync::Arc;
3532
3533        #[derive(Clone)]
3534        struct AppState {
3535            prefix: String,
3536        }
3537
3538        #[derive(Debug, Deserialize, JsonSchema)]
3539        struct QueryInput {
3540            query: String,
3541        }
3542
3543        let state = Arc::new(AppState {
3544            prefix: "db".to_string(),
3545        });
3546
3547        let tool = ToolBuilder::new("search")
3548            .description("Search")
3549            .extractor_handler(
3550                state,
3551                |State(app): State<Arc<AppState>>, Json(input): Json<QueryInput>| async move {
3552                    Ok(CallToolResult::text(format!(
3553                        "{}: {}",
3554                        app.prefix, input.query
3555                    )))
3556                },
3557            )
3558            .guard(|req: &ToolRequest| {
3559                let query = req.args.get("query").and_then(|v| v.as_str()).unwrap_or("");
3560                if query.is_empty() {
3561                    return Err("Query cannot be empty".to_string());
3562                }
3563                Ok(())
3564            })
3565            .build();
3566
3567        // Valid query
3568        let result = tool.call(serde_json::json!({"query": "hello"})).await;
3569        assert!(!result.is_error);
3570        assert_eq!(result.first_text().unwrap(), "db: hello");
3571
3572        // Empty query rejected by guard
3573        let result = tool.call(serde_json::json!({"query": ""})).await;
3574        assert!(result.is_error);
3575        assert!(
3576            result
3577                .first_text()
3578                .unwrap()
3579                .contains("Query cannot be empty")
3580        );
3581    }
3582
3583    #[tokio::test]
3584    async fn test_guard_on_extractor_handler_with_layer() {
3585        use std::sync::Arc;
3586        use std::time::Duration;
3587        use tower::timeout::TimeoutLayer;
3588
3589        #[derive(Clone)]
3590        struct AppState2 {
3591            prefix: String,
3592        }
3593
3594        #[derive(Debug, Deserialize, JsonSchema)]
3595        struct QueryInput2 {
3596            query: String,
3597        }
3598
3599        let state = Arc::new(AppState2 {
3600            prefix: "db".to_string(),
3601        });
3602
3603        let tool = ToolBuilder::new("search2")
3604            .description("Search with layer and guard")
3605            .extractor_handler(
3606                state,
3607                |State(app): State<Arc<AppState2>>, Json(input): Json<QueryInput2>| async move {
3608                    Ok(CallToolResult::text(format!(
3609                        "{}: {}",
3610                        app.prefix, input.query
3611                    )))
3612                },
3613            )
3614            .layer(TimeoutLayer::new(Duration::from_secs(5)))
3615            .guard(|_req: &ToolRequest| Ok(()))
3616            .build();
3617
3618        let result = tool.call(serde_json::json!({"query": "hello"})).await;
3619        assert!(!result.is_error);
3620        assert_eq!(result.first_text().unwrap(), "db: hello");
3621    }
3622
3623    #[tokio::test]
3624    async fn test_tool_with_guard_post_build() {
3625        let tool = ToolBuilder::new("admin_action")
3626            .description("Admin action")
3627            .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3628            .build();
3629
3630        // Apply guard after building
3631        let guarded = tool.with_guard(|req: &ToolRequest| {
3632            let name = req.args.get("name").and_then(|v| v.as_str()).unwrap_or("");
3633            if name == "admin" {
3634                Ok(())
3635            } else {
3636                Err("Only admin allowed".to_string())
3637            }
3638        });
3639
3640        // Admin passes
3641        let result = guarded.call(serde_json::json!({"name": "admin"})).await;
3642        assert!(!result.is_error);
3643
3644        // Non-admin blocked
3645        let result = guarded.call(serde_json::json!({"name": "user"})).await;
3646        assert!(result.is_error);
3647        assert!(result.first_text().unwrap().contains("Only admin allowed"));
3648    }
3649
3650    #[tokio::test]
3651    async fn test_with_guard_preserves_tool_metadata() {
3652        let tool = ToolBuilder::new("my_tool")
3653            .description("A tool")
3654            .title("My Tool")
3655            .read_only()
3656            .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3657            .build();
3658
3659        let guarded = tool.with_guard(|_req: &ToolRequest| Ok(()));
3660
3661        assert_eq!(guarded.name, "my_tool");
3662        assert_eq!(guarded.description.as_deref(), Some("A tool"));
3663        assert_eq!(guarded.title.as_deref(), Some("My Tool"));
3664        assert!(guarded.annotations.is_some());
3665    }
3666
3667    #[tokio::test]
3668    async fn test_guard_group_pattern() {
3669        // Demonstrate applying the same guard to multiple tools (per-group pattern)
3670        let require_auth = |req: &ToolRequest| {
3671            let token = req
3672                .args
3673                .get("_token")
3674                .and_then(|v| v.as_str())
3675                .unwrap_or("");
3676            if token == "valid" {
3677                Ok(())
3678            } else {
3679                Err("Authentication required".to_string())
3680            }
3681        };
3682
3683        let tool1 = ToolBuilder::new("action1")
3684            .description("Action 1")
3685            .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action1")) })
3686            .build();
3687        let tool2 = ToolBuilder::new("action2")
3688            .description("Action 2")
3689            .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action2")) })
3690            .build();
3691
3692        // Apply same guard to both
3693        let guarded1 = tool1.with_guard(require_auth);
3694        let guarded2 = tool2.with_guard(require_auth);
3695
3696        // Without auth
3697        let r1 = guarded1
3698            .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3699            .await;
3700        let r2 = guarded2
3701            .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3702            .await;
3703        assert!(r1.is_error);
3704        assert!(r2.is_error);
3705
3706        // With auth
3707        let r1 = guarded1
3708            .call(serde_json::json!({"name": "test", "_token": "valid"}))
3709            .await;
3710        let r2 = guarded2
3711            .call(serde_json::json!({"name": "test", "_token": "valid"}))
3712            .await;
3713        assert!(!r1.is_error);
3714        assert!(!r2.is_error);
3715    }
3716
3717    #[tokio::test]
3718    async fn test_input_validation_returns_tool_error() {
3719        // Per SEP-1303: input validation errors should be returned as
3720        // CallToolResult with isError=true, not as protocol errors.
3721        #[derive(Debug, Deserialize, JsonSchema)]
3722        struct StrictInput {
3723            name: String,
3724            count: u32,
3725        }
3726
3727        let tool = ToolBuilder::new("strict_tool")
3728            .description("requires specific input")
3729            .handler(|input: StrictInput| async move {
3730                Ok(CallToolResult::text(format!(
3731                    "{}: {}",
3732                    input.name, input.count
3733                )))
3734            })
3735            .build();
3736
3737        // Valid input works
3738        let result = tool
3739            .call(serde_json::json!({"name": "test", "count": 5}))
3740            .await;
3741        assert!(!result.is_error);
3742
3743        // Missing required field returns isError, not protocol error
3744        let result = tool.call(serde_json::json!({"name": "test"})).await;
3745        assert!(result.is_error);
3746        let text = result.first_text().unwrap();
3747        assert!(text.contains("Invalid input"), "got: {text}");
3748
3749        // Wrong type returns isError, not protocol error
3750        let result = tool
3751            .call(serde_json::json!({"name": "test", "count": "not_a_number"}))
3752            .await;
3753        assert!(result.is_error);
3754        let text = result.first_text().unwrap();
3755        assert!(text.contains("Invalid input"), "got: {text}");
3756    }
3757
3758    #[tokio::test]
3759    async fn test_input_schema_override_with_raw_args() {
3760        // With a RawArgs handler there is no typed input struct, so the
3761        // builder normally falls back to `{ "type": "object" }`. The
3762        // `input_schema` setter must let users declare a richer schema.
3763        let custom = serde_json::json!({
3764            "type": "object",
3765            "properties": {
3766                "query": { "type": "string", "minLength": 1 }
3767            },
3768            "required": ["query"]
3769        });
3770
3771        let tool = ToolBuilder::new("query")
3772            .description("Query with a custom schema")
3773            .input_schema(custom.clone())
3774            .extractor_handler((), |RawArgs(args): RawArgs| async move {
3775                Ok(CallToolResult::json(args))
3776            })
3777            .build();
3778
3779        let schema = tool.definition().input_schema;
3780        assert_eq!(schema, custom);
3781
3782        // The handler still executes against the raw args.
3783        let result = tool.call(serde_json::json!({"query": "hello"})).await;
3784        assert!(!result.is_error);
3785    }
3786
3787    #[tokio::test]
3788    async fn test_input_schema_override_wins_over_typed_handler() {
3789        // When both `.input_schema(...)` and a typed `.handler(|x: Foo|)` are
3790        // provided, the explicit schema must win over the schemars-generated
3791        // one.
3792        let custom = serde_json::json!({
3793            "type": "object",
3794            "title": "GreetOverride",
3795            "properties": {
3796                "name": { "type": "string", "minLength": 1, "maxLength": 64 }
3797            },
3798            "required": ["name"],
3799            "additionalProperties": false
3800        });
3801
3802        let tool = ToolBuilder::new("greet")
3803            .description("Greet someone with a hand-tuned schema")
3804            .input_schema(custom.clone())
3805            .handler(|input: GreetInput| async move {
3806                Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3807            })
3808            .build();
3809
3810        let schema = tool.definition().input_schema;
3811        assert_eq!(schema, custom);
3812        // Confirm the schemars-generated `GreetInput` schema did not leak in.
3813        assert_eq!(schema["title"], "GreetOverride");
3814
3815        // Handler still dispatches via the typed deserialization.
3816        let result = tool.call(serde_json::json!({"name": "World"})).await;
3817        assert!(!result.is_error);
3818    }
3819
3820    #[tokio::test]
3821    async fn test_input_schema_override_preserves_2020_12_constructs() {
3822        // Schemars cannot express `oneOf` in property positions directly;
3823        // overriding the schema must keep those advanced constructs intact.
3824        let custom = serde_json::json!({
3825            "type": "object",
3826            "properties": {
3827                "filter": {
3828                    "oneOf": [
3829                        { "type": "string" },
3830                        {
3831                            "type": "object",
3832                            "properties": { "field": { "type": "string" } },
3833                            "required": ["field"]
3834                        }
3835                    ]
3836                }
3837            },
3838            "required": ["filter"]
3839        });
3840
3841        let tool = ToolBuilder::new("filter_tool")
3842            .description("Demonstrates oneOf preservation")
3843            .input_schema(custom.clone())
3844            .extractor_handler((), |RawArgs(args): RawArgs| async move {
3845                Ok(CallToolResult::json(args))
3846            })
3847            .build();
3848
3849        let schema = tool.definition().input_schema;
3850        assert_eq!(schema, custom);
3851        let one_of = schema["properties"]["filter"]["oneOf"]
3852            .as_array()
3853            .expect("oneOf must survive as an array");
3854        assert_eq!(one_of.len(), 2);
3855        assert_eq!(one_of[0]["type"], "string");
3856        assert_eq!(one_of[1]["type"], "object");
3857    }
3858
3859    #[tokio::test]
3860    async fn test_input_schema_override_adds_type_object_if_missing() {
3861        // `ensure_object_schema` must still run against the user-supplied
3862        // schema, so MCP-spec `type: "object"` is added when omitted.
3863        let custom_no_type = serde_json::json!({
3864            "properties": {
3865                "x": { "type": "number" }
3866            }
3867        });
3868
3869        let tool = ToolBuilder::new("typeless")
3870            .description("Schema missing top-level type")
3871            .input_schema(custom_no_type)
3872            .extractor_handler((), |RawArgs(args): RawArgs| async move {
3873                Ok(CallToolResult::json(args))
3874            })
3875            .build();
3876
3877        let schema = tool.definition().input_schema;
3878        assert_eq!(schema["type"], "object");
3879        assert!(schema["properties"]["x"].is_object());
3880    }
3881}