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