tower_mcp/
extract.rs

1//! Extractor pattern for tool handlers
2//!
3//! This module provides an axum-inspired extractor pattern that makes state and context
4//! injection more declarative, reducing the combinatorial explosion of handler variants.
5//!
6//! # Overview
7//!
8//! Extractors implement [`FromToolRequest`], which extracts data from the tool request
9//! (context, state, and arguments). Multiple extractors can be combined in handler
10//! function parameters.
11//!
12//! # Built-in Extractors
13//!
14//! - [`Json<T>`] - Extract typed input from args (deserializes JSON)
15//! - [`State<T>`] - Extract shared state from per-tool state (cloned for each request)
16//! - [`Extension<T>`] - Extract data from router extensions (via `router.with_state()`)
17//! - [`Context`] - Extract the [`RequestContext`] for progress, cancellation, etc.
18//! - [`RawArgs`] - Extract raw `serde_json::Value` arguments
19//!
20//! ## State vs Extension
21//!
22//! - Use **`State<T>`** when state is passed directly to `extractor_handler()` (per-tool state)
23//! - Use **`Extension<T>`** when state is set via `McpRouter::with_state()` (router-level state)
24//!
25//! # schemars version alignment
26//!
27//! The `T` in [`Json<T>`] must implement [`schemars::JsonSchema`], and the
28//! derived impl must come from the same `schemars` major version tower-mcp uses
29//! (currently `1.x`). A version skew produces opaque `ExtractorHandler`
30//! trait-bound errors that do not name the real cause. Depend on `schemars`
31//! through the `tower_mcp::schemars` re-export to keep the versions aligned.
32//!
33//! # Example
34//!
35//! ```rust
36//! use std::sync::Arc;
37//! use tower_mcp::{ToolBuilder, CallToolResult};
38//! use tower_mcp::extract::{Json, State, Context};
39//! use schemars::JsonSchema;
40//! use serde::Deserialize;
41//!
42//! #[derive(Clone)]
43//! struct AppState {
44//!     db_url: String,
45//! }
46//!
47//! #[derive(Debug, Deserialize, JsonSchema)]
48//! struct QueryInput {
49//!     query: String,
50//! }
51//!
52//! let state = Arc::new(AppState { db_url: "postgres://...".to_string() });
53//!
54//! let tool = ToolBuilder::new("search")
55//!     .description("Search the database")
56//!     .extractor_handler(state, |
57//!         State(db): State<Arc<AppState>>,
58//!         ctx: Context,
59//!         Json(input): Json<QueryInput>,
60//!     | async move {
61//!         // Check cancellation
62//!         if ctx.is_cancelled() {
63//!             return Ok(CallToolResult::error("Cancelled"));
64//!         }
65//!         // Report progress
66//!         ctx.report_progress(0.5, Some(1.0), Some("Searching...")).await;
67//!         // Use state
68//!         Ok(CallToolResult::text(format!("Searched {} with query: {}", db.db_url, input.query)))
69//!     })
70//!     .build();
71//! ```
72//!
73//! # Extractor Order
74//!
75//! The order of extractors in the function signature doesn't matter. Each extractor
76//! independently extracts its data from the request.
77//!
78//! # Error Handling
79//!
80//! If an extractor fails (e.g., JSON deserialization fails), the handler returns
81//! a `CallToolResult::error()` with the rejection message.
82
83use std::future::Future;
84use std::marker::PhantomData;
85use std::ops::Deref;
86use std::pin::Pin;
87
88use schemars::JsonSchema;
89use serde::de::DeserializeOwned;
90use serde_json::Value;
91
92use crate::context::RequestContext;
93use crate::error::{Error, Result};
94use crate::protocol::CallToolResult;
95
96// =============================================================================
97// Rejection Types
98// =============================================================================
99
100/// A simple rejection with a message string.
101///
102/// This is a general-purpose rejection type for custom extractors.
103/// For more specific error information, use the typed rejection types
104/// like [`JsonRejection`] or [`ExtensionRejection`].
105#[derive(Debug, Clone)]
106pub struct Rejection {
107    message: String,
108}
109
110impl Rejection {
111    /// Create a new rejection with the given message.
112    pub fn new(message: impl Into<String>) -> Self {
113        Self {
114            message: message.into(),
115        }
116    }
117
118    /// Get the rejection message.
119    pub fn message(&self) -> &str {
120        &self.message
121    }
122}
123
124impl std::fmt::Display for Rejection {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        write!(f, "{}", self.message)
127    }
128}
129
130impl std::error::Error for Rejection {}
131
132impl From<Rejection> for Error {
133    fn from(rejection: Rejection) -> Self {
134        Error::tool(rejection.message)
135    }
136}
137
138/// Rejection returned when JSON deserialization fails.
139///
140/// This rejection provides structured information about the deserialization
141/// error, including the path to the failing field when available.
142///
143/// # Example
144///
145/// ```rust
146/// use tower_mcp::extract::JsonRejection;
147///
148/// let rejection = JsonRejection::new("missing field `name`");
149/// assert!(rejection.message().contains("name"));
150/// ```
151#[derive(Debug, Clone)]
152pub struct JsonRejection {
153    message: String,
154    /// The serde error path, if available (e.g., "users[0].name")
155    path: Option<String>,
156}
157
158impl JsonRejection {
159    /// Create a new JSON rejection from a serde error.
160    pub fn new(message: impl Into<String>) -> Self {
161        Self {
162            message: message.into(),
163            path: None,
164        }
165    }
166
167    /// Create a JSON rejection with a path to the failing field.
168    pub fn with_path(message: impl Into<String>, path: impl Into<String>) -> Self {
169        Self {
170            message: message.into(),
171            path: Some(path.into()),
172        }
173    }
174
175    /// Get the error message.
176    pub fn message(&self) -> &str {
177        &self.message
178    }
179
180    /// Get the path to the failing field, if available.
181    pub fn path(&self) -> Option<&str> {
182        self.path.as_deref()
183    }
184}
185
186impl std::fmt::Display for JsonRejection {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        if let Some(path) = &self.path {
189            write!(f, "Invalid input at `{}`: {}", path, self.message)
190        } else {
191            write!(f, "Invalid input: {}", self.message)
192        }
193    }
194}
195
196impl std::error::Error for JsonRejection {}
197
198impl From<JsonRejection> for Error {
199    fn from(rejection: JsonRejection) -> Self {
200        Error::tool(rejection.to_string())
201    }
202}
203
204impl From<serde_json::Error> for JsonRejection {
205    fn from(err: serde_json::Error) -> Self {
206        // Try to extract path information from serde error
207        let path = if err.is_data() {
208            // serde_json provides line/column but not field path in the error itself
209            // The path is embedded in the message for some error types
210            None
211        } else {
212            None
213        };
214
215        Self {
216            message: err.to_string(),
217            path,
218        }
219    }
220}
221
222/// Rejection returned when an extension is not found.
223///
224/// This rejection is returned by the [`Extension`] extractor when the
225/// requested type is not present in the router's extensions.
226///
227/// # Example
228///
229/// ```rust
230/// use tower_mcp::extract::ExtensionRejection;
231///
232/// let rejection = ExtensionRejection::not_found::<String>();
233/// assert!(rejection.type_name().contains("String"));
234/// ```
235#[derive(Debug, Clone)]
236pub struct ExtensionRejection {
237    type_name: &'static str,
238}
239
240impl ExtensionRejection {
241    /// Create a rejection for a missing extension type.
242    pub fn not_found<T>() -> Self {
243        Self {
244            type_name: std::any::type_name::<T>(),
245        }
246    }
247
248    /// Get the type name of the missing extension.
249    pub fn type_name(&self) -> &'static str {
250        self.type_name
251    }
252}
253
254impl std::fmt::Display for ExtensionRejection {
255    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
256        write!(
257            f,
258            "Extension of type `{}` not found. Did you call `router.with_state()` or `router.with_extension()`?",
259            self.type_name
260        )
261    }
262}
263
264impl std::error::Error for ExtensionRejection {}
265
266impl From<ExtensionRejection> for Error {
267    fn from(rejection: ExtensionRejection) -> Self {
268        Error::tool(rejection.to_string())
269    }
270}
271
272/// Trait for extracting data from a tool request.
273///
274/// Implement this trait to create custom extractors that can be used
275/// in `extractor_handler` functions.
276///
277/// # Type Parameters
278///
279/// - `S` - The state type. Defaults to `()` for extractors that don't need state.
280///
281/// # Example
282///
283/// ```rust
284/// use tower_mcp::extract::{FromToolRequest, Rejection};
285/// use tower_mcp::RequestContext;
286/// use serde_json::Value;
287///
288/// struct RequestId(String);
289///
290/// impl<S> FromToolRequest<S> for RequestId {
291///     type Rejection = Rejection;
292///
293///     fn from_tool_request(
294///         ctx: &RequestContext,
295///         _state: &S,
296///         _args: &Value,
297///     ) -> Result<Self, Self::Rejection> {
298///         Ok(RequestId(format!("{:?}", ctx.request_id())))
299///     }
300/// }
301/// ```
302pub trait FromToolRequest<S = ()>: Sized {
303    /// The rejection type returned when extraction fails.
304    type Rejection: Into<Error>;
305
306    /// Extract this type from the tool request.
307    ///
308    /// # Arguments
309    ///
310    /// * `ctx` - The request context with progress, cancellation, etc.
311    /// * `state` - The shared state passed to the handler
312    /// * `args` - The raw JSON arguments to the tool
313    fn from_tool_request(
314        ctx: &RequestContext,
315        state: &S,
316        args: &Value,
317    ) -> std::result::Result<Self, Self::Rejection>;
318}
319
320// =============================================================================
321// Built-in Extractors
322// =============================================================================
323
324/// Extract and deserialize JSON arguments into a typed struct.
325///
326/// This extractor deserializes the tool's JSON arguments into type `T`.
327/// The type must implement [`serde::de::DeserializeOwned`] and [`schemars::JsonSchema`].
328///
329/// # Example
330///
331/// ```rust
332/// use tower_mcp::extract::Json;
333/// use schemars::JsonSchema;
334/// use serde::Deserialize;
335///
336/// #[derive(Debug, Deserialize, JsonSchema)]
337/// struct MyInput {
338///     name: String,
339///     count: i32,
340/// }
341///
342/// // In an extractor handler:
343/// // |Json(input): Json<MyInput>| async move { ... }
344/// ```
345///
346/// # Rejection
347///
348/// Returns a [`JsonRejection`] if deserialization fails. The rejection contains
349/// the error message and potentially the path to the failing field.
350#[derive(Debug, Clone, Copy)]
351pub struct Json<T>(pub T);
352
353impl<T> Deref for Json<T> {
354    type Target = T;
355
356    fn deref(&self) -> &Self::Target {
357        &self.0
358    }
359}
360
361impl<S, T> FromToolRequest<S> for Json<T>
362where
363    T: DeserializeOwned,
364{
365    type Rejection = JsonRejection;
366
367    fn from_tool_request(
368        _ctx: &RequestContext,
369        _state: &S,
370        args: &Value,
371    ) -> std::result::Result<Self, Self::Rejection> {
372        serde_json::from_value(args.clone())
373            .map(Json)
374            .map_err(JsonRejection::from)
375    }
376}
377
378/// Extract shared state.
379///
380/// This extractor clones the state passed to `extractor_handler` and provides
381/// it to the handler. The state type must match the type passed to the builder.
382///
383/// # Example
384///
385/// ```rust
386/// use std::sync::Arc;
387/// use tower_mcp::extract::State;
388///
389/// #[derive(Clone)]
390/// struct AppState {
391///     db_url: String,
392/// }
393///
394/// // In an extractor handler:
395/// // |State(state): State<Arc<AppState>>| async move { ... }
396/// ```
397///
398/// # Note
399///
400/// For expensive-to-clone types, wrap them in `Arc` before passing to
401/// `extractor_handler`.
402#[derive(Debug, Clone, Copy)]
403pub struct State<T>(pub T);
404
405impl<T> Deref for State<T> {
406    type Target = T;
407
408    fn deref(&self) -> &Self::Target {
409        &self.0
410    }
411}
412
413impl<S: Clone> FromToolRequest<S> for State<S> {
414    type Rejection = Rejection;
415
416    fn from_tool_request(
417        _ctx: &RequestContext,
418        state: &S,
419        _args: &Value,
420    ) -> std::result::Result<Self, Self::Rejection> {
421        Ok(State(state.clone()))
422    }
423}
424
425/// Extract the request context.
426///
427/// This extractor provides access to the [`RequestContext`], which contains:
428/// - Progress reporting via `report_progress()`
429/// - Cancellation checking via `is_cancelled()`
430/// - Sampling capabilities via `sample()`
431/// - Elicitation capabilities via `elicit_form()` and `elicit_url()`
432/// - Log sending via `send_log()`
433///
434/// # Example
435///
436/// ```rust
437/// use tower_mcp::extract::Context;
438///
439/// // In an extractor handler:
440/// // |ctx: Context| async move {
441/// //     ctx.report_progress(0.5, Some(1.0), Some("Working...")).await;
442/// //     // ...
443/// // }
444/// ```
445#[derive(Debug, Clone)]
446pub struct Context(RequestContext);
447
448impl Context {
449    /// Get the inner RequestContext
450    pub fn into_inner(self) -> RequestContext {
451        self.0
452    }
453}
454
455impl Deref for Context {
456    type Target = RequestContext;
457
458    fn deref(&self) -> &Self::Target {
459        &self.0
460    }
461}
462
463impl<S> FromToolRequest<S> for Context {
464    type Rejection = Rejection;
465
466    fn from_tool_request(
467        ctx: &RequestContext,
468        _state: &S,
469        _args: &Value,
470    ) -> std::result::Result<Self, Self::Rejection> {
471        Ok(Context(ctx.clone()))
472    }
473}
474
475/// Extract raw JSON arguments.
476///
477/// This extractor provides the raw `serde_json::Value` arguments without
478/// any deserialization. Useful when you need full control over argument
479/// parsing or when the schema is dynamic.
480///
481/// # Example
482///
483/// ```rust
484/// use tower_mcp::extract::RawArgs;
485///
486/// // In an extractor handler:
487/// // |RawArgs(args): RawArgs| async move {
488/// //     // args is serde_json::Value
489/// //     if let Some(name) = args.get("name") { ... }
490/// // }
491/// ```
492#[derive(Debug, Clone)]
493pub struct RawArgs(pub Value);
494
495impl Deref for RawArgs {
496    type Target = Value;
497
498    fn deref(&self) -> &Self::Target {
499        &self.0
500    }
501}
502
503impl<S> FromToolRequest<S> for RawArgs {
504    type Rejection = Rejection;
505
506    fn from_tool_request(
507        _ctx: &RequestContext,
508        _state: &S,
509        args: &Value,
510    ) -> std::result::Result<Self, Self::Rejection> {
511        Ok(RawArgs(args.clone()))
512    }
513}
514
515/// Extract typed data from router extensions.
516///
517/// This extractor retrieves data that was added to the router via
518/// [`crate::McpRouter::with_state()`] or [`crate::McpRouter::with_extension()`], or
519/// inserted by middleware into the request context's extensions.
520///
521/// # Example
522///
523/// ```rust
524/// use std::sync::Arc;
525/// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
526/// use tower_mcp::extract::{Extension, Json};
527/// use schemars::JsonSchema;
528/// use serde::Deserialize;
529///
530/// #[derive(Clone)]
531/// struct DatabasePool {
532///     url: String,
533/// }
534///
535/// #[derive(Deserialize, JsonSchema)]
536/// struct QueryInput {
537///     sql: String,
538/// }
539///
540/// let pool = Arc::new(DatabasePool { url: "postgres://...".into() });
541///
542/// let tool = ToolBuilder::new("query")
543///     .description("Run a query")
544///     .extractor_handler(
545///         (),
546///         |Extension(db): Extension<Arc<DatabasePool>>, Json(input): Json<QueryInput>| async move {
547///             Ok(CallToolResult::text(format!("Query on {}: {}", db.url, input.sql)))
548///         },
549///     )
550///     .build();
551///
552/// let router = McpRouter::new()
553///     .with_state(pool)
554///     .tool(tool);
555/// ```
556///
557/// # Rejection
558///
559/// Returns an [`ExtensionRejection`] if the requested type is not found in the extensions.
560/// The rejection contains the type name of the missing extension.
561#[derive(Debug, Clone)]
562pub struct Extension<T>(pub T);
563
564impl<T> Deref for Extension<T> {
565    type Target = T;
566
567    fn deref(&self) -> &Self::Target {
568        &self.0
569    }
570}
571
572impl<S, T> FromToolRequest<S> for Extension<T>
573where
574    T: Clone + Send + Sync + 'static,
575{
576    type Rejection = ExtensionRejection;
577
578    fn from_tool_request(
579        ctx: &RequestContext,
580        _state: &S,
581        _args: &Value,
582    ) -> std::result::Result<Self, Self::Rejection> {
583        ctx.extension::<T>()
584            .cloned()
585            .map(Extension)
586            .ok_or_else(ExtensionRejection::not_found::<T>)
587    }
588}
589
590// =============================================================================
591// Handler Trait
592// =============================================================================
593
594/// A handler that uses extractors.
595///
596/// This trait is implemented for functions that take extractors as arguments.
597/// You don't need to implement this trait directly; it's automatically
598/// implemented for compatible async functions.
599#[diagnostic::on_unimplemented(
600    message = "`{Self}` is not a valid extractor handler",
601    note = "each closure argument must be an extractor (`Json<T>`, `State<S>`, `Context`, `Extension<T>`, `RawArgs`) and the return type must be `Result<impl Into<CallToolResult>, ToolError>`",
602    note = "for a `Json<T>` argument, `T` must implement `serde::Deserialize` and `schemars::JsonSchema`",
603    note = "if `T` derives `JsonSchema` but this still fails, check for a `schemars` major-version mismatch: the derive must come from the same `schemars` version tower-mcp uses (>=1). Depend on it via the `tower_mcp::schemars` re-export to stay aligned"
604)]
605pub trait ExtractorHandler<S, T>: Clone + Send + Sync + 'static {
606    /// The future returned by the handler.
607    type Future: Future<Output = Result<CallToolResult>> + Send;
608
609    /// Call the handler with extracted values.
610    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future;
611
612    /// Get the input schema for this handler.
613    ///
614    /// Returns `None` if no `Json<T>` extractor is used.
615    fn input_schema() -> Value;
616}
617
618// Implementation for single extractor
619impl<S, F, Fut, T1> ExtractorHandler<S, (T1,)> for F
620where
621    S: Clone + Send + Sync + 'static,
622    F: Fn(T1) -> Fut + Clone + Send + Sync + 'static,
623    Fut: Future<Output = Result<CallToolResult>> + Send,
624    T1: FromToolRequest<S> + HasSchema + Send,
625{
626    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
627
628    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
629        Box::pin(async move {
630            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
631            self(t1).await
632        })
633    }
634
635    fn input_schema() -> Value {
636        if let Some(schema) = T1::schema() {
637            return schema;
638        }
639        serde_json::json!({
640            "type": "object",
641            "additionalProperties": true
642        })
643    }
644}
645
646// Implementation for two extractors
647impl<S, F, Fut, T1, T2> ExtractorHandler<S, (T1, T2)> for F
648where
649    S: Clone + Send + Sync + 'static,
650    F: Fn(T1, T2) -> Fut + Clone + Send + Sync + 'static,
651    Fut: Future<Output = Result<CallToolResult>> + Send,
652    T1: FromToolRequest<S> + HasSchema + Send,
653    T2: FromToolRequest<S> + HasSchema + Send,
654{
655    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
656
657    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
658        Box::pin(async move {
659            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
660            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
661            self(t1, t2).await
662        })
663    }
664
665    fn input_schema() -> Value {
666        if let Some(schema) = T2::schema() {
667            return schema;
668        }
669        if let Some(schema) = T1::schema() {
670            return schema;
671        }
672        serde_json::json!({
673            "type": "object",
674            "additionalProperties": true
675        })
676    }
677}
678
679// Implementation for three extractors
680impl<S, F, Fut, T1, T2, T3> ExtractorHandler<S, (T1, T2, T3)> for F
681where
682    S: Clone + Send + Sync + 'static,
683    F: Fn(T1, T2, T3) -> Fut + Clone + Send + Sync + 'static,
684    Fut: Future<Output = Result<CallToolResult>> + Send,
685    T1: FromToolRequest<S> + HasSchema + Send,
686    T2: FromToolRequest<S> + HasSchema + Send,
687    T3: FromToolRequest<S> + HasSchema + Send,
688{
689    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
690
691    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
692        Box::pin(async move {
693            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
694            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
695            let t3 = T3::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
696            self(t1, t2, t3).await
697        })
698    }
699
700    fn input_schema() -> Value {
701        if let Some(schema) = T3::schema() {
702            return schema;
703        }
704        if let Some(schema) = T2::schema() {
705            return schema;
706        }
707        if let Some(schema) = T1::schema() {
708            return schema;
709        }
710        serde_json::json!({
711            "type": "object",
712            "additionalProperties": true
713        })
714    }
715}
716
717// Implementation for four extractors
718impl<S, F, Fut, T1, T2, T3, T4> ExtractorHandler<S, (T1, T2, T3, T4)> for F
719where
720    S: Clone + Send + Sync + 'static,
721    F: Fn(T1, T2, T3, T4) -> Fut + Clone + Send + Sync + 'static,
722    Fut: Future<Output = Result<CallToolResult>> + Send,
723    T1: FromToolRequest<S> + HasSchema + Send,
724    T2: FromToolRequest<S> + HasSchema + Send,
725    T3: FromToolRequest<S> + HasSchema + Send,
726    T4: FromToolRequest<S> + HasSchema + Send,
727{
728    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
729
730    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
731        Box::pin(async move {
732            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
733            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
734            let t3 = T3::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
735            let t4 = T4::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
736            self(t1, t2, t3, t4).await
737        })
738    }
739
740    fn input_schema() -> Value {
741        if let Some(schema) = T4::schema() {
742            return schema;
743        }
744        if let Some(schema) = T3::schema() {
745            return schema;
746        }
747        if let Some(schema) = T2::schema() {
748            return schema;
749        }
750        if let Some(schema) = T1::schema() {
751            return schema;
752        }
753        serde_json::json!({
754            "type": "object",
755            "additionalProperties": true
756        })
757    }
758}
759
760// Implementation for five extractors
761impl<S, F, Fut, T1, T2, T3, T4, T5> ExtractorHandler<S, (T1, T2, T3, T4, T5)> for F
762where
763    S: Clone + Send + Sync + 'static,
764    F: Fn(T1, T2, T3, T4, T5) -> Fut + Clone + Send + Sync + 'static,
765    Fut: Future<Output = Result<CallToolResult>> + Send,
766    T1: FromToolRequest<S> + HasSchema + Send,
767    T2: FromToolRequest<S> + HasSchema + Send,
768    T3: FromToolRequest<S> + HasSchema + Send,
769    T4: FromToolRequest<S> + HasSchema + Send,
770    T5: FromToolRequest<S> + HasSchema + Send,
771{
772    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
773
774    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
775        Box::pin(async move {
776            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
777            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
778            let t3 = T3::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
779            let t4 = T4::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
780            let t5 = T5::from_tool_request(&ctx, &state, &args).map_err(Into::into)?;
781            self(t1, t2, t3, t4, t5).await
782        })
783    }
784
785    fn input_schema() -> Value {
786        if let Some(schema) = T5::schema() {
787            return schema;
788        }
789        if let Some(schema) = T4::schema() {
790            return schema;
791        }
792        if let Some(schema) = T3::schema() {
793            return schema;
794        }
795        if let Some(schema) = T2::schema() {
796            return schema;
797        }
798        if let Some(schema) = T1::schema() {
799            return schema;
800        }
801        serde_json::json!({
802            "type": "object",
803            "additionalProperties": true
804        })
805    }
806}
807
808// =============================================================================
809// Schema Extraction Helper
810// =============================================================================
811
812/// Helper trait to get schema from `Json<T>` extractor
813#[diagnostic::on_unimplemented(
814    message = "`{Self}` does not implement `HasSchema`",
815    note = "for `Json<T>` this means `T: schemars::JsonSchema` is not satisfied",
816    note = "a common cause is a `schemars` major-version mismatch: the derive on `T` must come from the same `schemars` version tower-mcp uses (>=1)",
817    note = "depend on `schemars` via the `tower_mcp::schemars` re-export to keep the versions aligned"
818)]
819pub trait HasSchema {
820    /// Returns the JSON Schema for this type, if available.
821    fn schema() -> Option<Value>;
822}
823
824impl<T: JsonSchema> HasSchema for Json<T> {
825    fn schema() -> Option<Value> {
826        let schema = schemars::schema_for!(T);
827        serde_json::to_value(schema)
828            .ok()
829            .map(crate::tool::ensure_object_schema)
830    }
831}
832
833// Default impl for non-Json extractors
834impl HasSchema for Context {
835    fn schema() -> Option<Value> {
836        None
837    }
838}
839
840impl HasSchema for RawArgs {
841    fn schema() -> Option<Value> {
842        None
843    }
844}
845
846impl<T> HasSchema for State<T> {
847    fn schema() -> Option<Value> {
848        None
849    }
850}
851
852impl<T> HasSchema for Extension<T> {
853    fn schema() -> Option<Value> {
854        None
855    }
856}
857
858// =============================================================================
859// Typed Extractor Handler
860// =============================================================================
861
862/// A handler that uses extractors with typed JSON input.
863///
864/// This trait is similar to [`ExtractorHandler`] but provides proper JSON
865/// schema generation for the input type when `Json<T>` is used.
866#[deprecated(
867    since = "0.8.0",
868    note = "Use `ExtractorHandler` instead -- `extractor_handler` auto-detects JSON schema from `Json<T>` extractors"
869)]
870pub trait TypedExtractorHandler<S, T, I>: Clone + Send + Sync + 'static
871where
872    I: JsonSchema,
873{
874    /// The future returned by the handler.
875    type Future: Future<Output = Result<CallToolResult>> + Send;
876
877    /// Call the handler with extracted values.
878    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future;
879}
880
881// Single extractor with Json<T>
882#[allow(deprecated)]
883impl<S, F, Fut, T> TypedExtractorHandler<S, (Json<T>,), T> for F
884where
885    S: Clone + Send + Sync + 'static,
886    F: Fn(Json<T>) -> Fut + Clone + Send + Sync + 'static,
887    Fut: Future<Output = Result<CallToolResult>> + Send,
888    T: DeserializeOwned + JsonSchema + Send,
889{
890    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
891
892    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
893        Box::pin(async move {
894            let t1 =
895                Json::<T>::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
896            self(t1).await
897        })
898    }
899}
900
901// Two extractors ending with Json<T>
902#[allow(deprecated)]
903impl<S, F, Fut, T1, T> TypedExtractorHandler<S, (T1, Json<T>), T> for F
904where
905    S: Clone + Send + Sync + 'static,
906    F: Fn(T1, Json<T>) -> Fut + Clone + Send + Sync + 'static,
907    Fut: Future<Output = Result<CallToolResult>> + Send,
908    T1: FromToolRequest<S> + Send,
909    T: DeserializeOwned + JsonSchema + Send,
910{
911    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
912
913    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
914        Box::pin(async move {
915            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
916            let t2 =
917                Json::<T>::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
918            self(t1, t2).await
919        })
920    }
921}
922
923// Three extractors ending with Json<T>
924#[allow(deprecated)]
925impl<S, F, Fut, T1, T2, T> TypedExtractorHandler<S, (T1, T2, Json<T>), T> for F
926where
927    S: Clone + Send + Sync + 'static,
928    F: Fn(T1, T2, Json<T>) -> Fut + Clone + Send + Sync + 'static,
929    Fut: Future<Output = Result<CallToolResult>> + Send,
930    T1: FromToolRequest<S> + Send,
931    T2: FromToolRequest<S> + Send,
932    T: DeserializeOwned + JsonSchema + Send,
933{
934    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
935
936    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
937        Box::pin(async move {
938            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
939            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
940            let t3 =
941                Json::<T>::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
942            self(t1, t2, t3).await
943        })
944    }
945}
946
947// Four extractors ending with Json<T>
948#[allow(deprecated)]
949impl<S, F, Fut, T1, T2, T3, T> TypedExtractorHandler<S, (T1, T2, T3, Json<T>), T> for F
950where
951    S: Clone + Send + Sync + 'static,
952    F: Fn(T1, T2, T3, Json<T>) -> Fut + Clone + Send + Sync + 'static,
953    Fut: Future<Output = Result<CallToolResult>> + Send,
954    T1: FromToolRequest<S> + Send,
955    T2: FromToolRequest<S> + Send,
956    T3: FromToolRequest<S> + Send,
957    T: DeserializeOwned + JsonSchema + Send,
958{
959    type Future = Pin<Box<dyn Future<Output = Result<CallToolResult>> + Send>>;
960
961    fn call(self, ctx: RequestContext, state: S, args: Value) -> Self::Future {
962        Box::pin(async move {
963            let t1 = T1::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
964            let t2 = T2::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
965            let t3 = T3::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
966            let t4 =
967                Json::<T>::from_tool_request(&ctx, &state, &args).map_err(Into::<Error>::into)?;
968            self(t1, t2, t3, t4).await
969        })
970    }
971}
972
973// =============================================================================
974// ToolBuilder Extensions
975// =============================================================================
976
977use crate::tool::{
978    BoxFuture, GuardLayer, Tool, ToolCatchError, ToolHandler, ToolHandlerService, ToolRequest,
979};
980use tower::util::BoxCloneService;
981use tower_service::Service;
982
983/// Internal handler wrapper for extractor-based handlers
984pub(crate) struct ExtractorToolHandler<S, F, T> {
985    state: S,
986    handler: F,
987    input_schema: Value,
988    _phantom: PhantomData<T>,
989}
990
991impl<S, F, T> ToolHandler for ExtractorToolHandler<S, F, T>
992where
993    S: Clone + Send + Sync + 'static,
994    F: ExtractorHandler<S, T> + Clone,
995    T: Send + Sync + 'static,
996{
997    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
998        let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
999        self.call_with_context(ctx, args)
1000    }
1001
1002    fn call_with_context(
1003        &self,
1004        ctx: RequestContext,
1005        args: Value,
1006    ) -> BoxFuture<'_, Result<CallToolResult>> {
1007        let state = self.state.clone();
1008        let handler = self.handler.clone();
1009        Box::pin(async move { handler.call(ctx, state, args).await })
1010    }
1011
1012    fn uses_context(&self) -> bool {
1013        true
1014    }
1015
1016    fn input_schema(&self) -> Value {
1017        self.input_schema.clone()
1018    }
1019}
1020
1021/// Builder state for extractor-based handlers
1022#[doc(hidden)]
1023pub struct ToolBuilderWithExtractor<S, F, T> {
1024    pub(crate) name: String,
1025    pub(crate) title: Option<String>,
1026    pub(crate) description: Option<String>,
1027    pub(crate) output_schema: Option<Value>,
1028    pub(crate) icons: Option<Vec<crate::protocol::ToolIcon>>,
1029    pub(crate) annotations: Option<crate::protocol::ToolAnnotations>,
1030    pub(crate) task_support: crate::protocol::TaskSupportMode,
1031    pub(crate) state: S,
1032    pub(crate) handler: F,
1033    pub(crate) input_schema: Value,
1034    pub(crate) _phantom: PhantomData<T>,
1035}
1036
1037impl<S, F, T> ToolBuilderWithExtractor<S, F, T>
1038where
1039    S: Clone + Send + Sync + 'static,
1040    F: ExtractorHandler<S, T> + Clone,
1041    T: Send + Sync + 'static,
1042{
1043    /// Build the tool.
1044    pub fn build(self) -> Tool {
1045        let handler = ExtractorToolHandler {
1046            state: self.state,
1047            handler: self.handler,
1048            input_schema: self.input_schema.clone(),
1049            _phantom: PhantomData,
1050        };
1051
1052        let handler_service = ToolHandlerService::new(handler);
1053        let catch_error = ToolCatchError::new(handler_service);
1054        let service = BoxCloneService::new(catch_error);
1055
1056        Tool {
1057            name: self.name,
1058            title: self.title,
1059            description: self.description,
1060            output_schema: self.output_schema,
1061            icons: self.icons,
1062            annotations: self.annotations,
1063            meta: None,
1064            task_support: self.task_support,
1065            required_client_capabilities: None,
1066            service: Some(service),
1067            #[cfg(feature = "stateless")]
1068            mrtr_handler: None,
1069            input_schema: self.input_schema,
1070        }
1071    }
1072
1073    /// Apply a Tower layer (middleware) to this tool.
1074    ///
1075    /// The layer wraps the tool's handler service, enabling functionality like
1076    /// timeouts, rate limiting, and metrics collection at the per-tool level.
1077    ///
1078    /// # Example
1079    ///
1080    /// ```rust
1081    /// use std::sync::Arc;
1082    /// use std::time::Duration;
1083    /// use tower::timeout::TimeoutLayer;
1084    /// use tower_mcp::{ToolBuilder, CallToolResult};
1085    /// use tower_mcp::extract::{Json, State};
1086    /// use schemars::JsonSchema;
1087    /// use serde::Deserialize;
1088    ///
1089    /// #[derive(Clone)]
1090    /// struct AppState { prefix: String }
1091    ///
1092    /// #[derive(Debug, Deserialize, JsonSchema)]
1093    /// struct QueryInput { query: String }
1094    ///
1095    /// let state = Arc::new(AppState { prefix: "db".to_string() });
1096    ///
1097    /// let tool = ToolBuilder::new("search")
1098    ///     .description("Search with timeout")
1099    ///     .extractor_handler(state, |
1100    ///         State(app): State<Arc<AppState>>,
1101    ///         Json(input): Json<QueryInput>,
1102    ///     | async move {
1103    ///         Ok(CallToolResult::text(format!("{}: {}", app.prefix, input.query)))
1104    ///     })
1105    ///     .layer(TimeoutLayer::new(Duration::from_secs(30)))
1106    ///     .build();
1107    /// ```
1108    pub fn layer<L>(self, layer: L) -> ToolBuilderWithExtractorLayer<S, F, T, L> {
1109        ToolBuilderWithExtractorLayer {
1110            name: self.name,
1111            title: self.title,
1112            description: self.description,
1113            output_schema: self.output_schema,
1114            icons: self.icons,
1115            annotations: self.annotations,
1116            task_support: self.task_support,
1117            state: self.state,
1118            handler: self.handler,
1119            input_schema: self.input_schema,
1120            layer,
1121            _phantom: PhantomData,
1122        }
1123    }
1124
1125    /// Apply a guard to this tool.
1126    ///
1127    /// See [`ToolBuilderWithHandler::guard`](crate::ToolBuilder) for details.
1128    pub fn guard<G>(self, guard: G) -> ToolBuilderWithExtractorLayer<S, F, T, GuardLayer<G>>
1129    where
1130        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1131    {
1132        self.layer(GuardLayer::new(guard))
1133    }
1134}
1135
1136/// Builder state after a layer has been applied to an extractor handler.
1137///
1138/// This builder allows chaining additional layers and building the final tool.
1139#[doc(hidden)]
1140pub struct ToolBuilderWithExtractorLayer<S, F, T, L> {
1141    name: String,
1142    title: Option<String>,
1143    description: Option<String>,
1144    output_schema: Option<Value>,
1145    icons: Option<Vec<crate::protocol::ToolIcon>>,
1146    annotations: Option<crate::protocol::ToolAnnotations>,
1147    task_support: crate::protocol::TaskSupportMode,
1148    state: S,
1149    handler: F,
1150    input_schema: Value,
1151    layer: L,
1152    _phantom: PhantomData<T>,
1153}
1154
1155#[allow(private_bounds)]
1156impl<S, F, T, L> ToolBuilderWithExtractorLayer<S, F, T, L>
1157where
1158    S: Clone + Send + Sync + 'static,
1159    F: ExtractorHandler<S, T> + Clone,
1160    T: Send + Sync + 'static,
1161    L: tower::Layer<ToolHandlerService<ExtractorToolHandler<S, F, T>>>
1162        + Clone
1163        + Send
1164        + Sync
1165        + 'static,
1166    L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
1167    <L::Service as Service<ToolRequest>>::Error: std::fmt::Display + Send,
1168    <L::Service as Service<ToolRequest>>::Future: Send,
1169{
1170    /// Build the tool with the applied layer(s).
1171    pub fn build(self) -> Tool {
1172        let handler = ExtractorToolHandler {
1173            state: self.state,
1174            handler: self.handler,
1175            input_schema: self.input_schema.clone(),
1176            _phantom: PhantomData,
1177        };
1178
1179        let handler_service = ToolHandlerService::new(handler);
1180        let layered = self.layer.layer(handler_service);
1181        let catch_error = ToolCatchError::new(layered);
1182        let service = BoxCloneService::new(catch_error);
1183
1184        Tool {
1185            name: self.name,
1186            title: self.title,
1187            description: self.description,
1188            output_schema: self.output_schema,
1189            icons: self.icons,
1190            annotations: self.annotations,
1191            meta: None,
1192            task_support: self.task_support,
1193            required_client_capabilities: None,
1194            service: Some(service),
1195            #[cfg(feature = "stateless")]
1196            mrtr_handler: None,
1197            input_schema: self.input_schema,
1198        }
1199    }
1200
1201    /// Apply an additional Tower layer (middleware).
1202    ///
1203    /// Layers are applied in order, with earlier layers wrapping later ones.
1204    /// This means the first layer added is the outermost middleware.
1205    pub fn layer<L2>(
1206        self,
1207        layer: L2,
1208    ) -> ToolBuilderWithExtractorLayer<S, F, T, tower::layer::util::Stack<L2, L>> {
1209        ToolBuilderWithExtractorLayer {
1210            name: self.name,
1211            title: self.title,
1212            description: self.description,
1213            output_schema: self.output_schema,
1214            icons: self.icons,
1215            annotations: self.annotations,
1216            task_support: self.task_support,
1217            state: self.state,
1218            handler: self.handler,
1219            input_schema: self.input_schema,
1220            layer: tower::layer::util::Stack::new(layer, self.layer),
1221            _phantom: PhantomData,
1222        }
1223    }
1224
1225    /// Apply a guard to this tool.
1226    ///
1227    /// See [`ToolBuilderWithHandler::guard`](crate::ToolBuilder) for details.
1228    pub fn guard<G>(
1229        self,
1230        guard: G,
1231    ) -> ToolBuilderWithExtractorLayer<S, F, T, tower::layer::util::Stack<GuardLayer<G>, L>>
1232    where
1233        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1234    {
1235        self.layer(GuardLayer::new(guard))
1236    }
1237}
1238
1239/// Builder state for extractor-based handlers with typed JSON input
1240#[doc(hidden)]
1241#[deprecated(
1242    since = "0.8.0",
1243    note = "Use `ToolBuilderWithExtractor` via `extractor_handler` instead"
1244)]
1245pub struct ToolBuilderWithTypedExtractor<S, F, T, I> {
1246    pub(crate) name: String,
1247    pub(crate) title: Option<String>,
1248    pub(crate) description: Option<String>,
1249    pub(crate) output_schema: Option<Value>,
1250    pub(crate) input_schema_override: Option<Value>,
1251    pub(crate) icons: Option<Vec<crate::protocol::ToolIcon>>,
1252    pub(crate) annotations: Option<crate::protocol::ToolAnnotations>,
1253    pub(crate) task_support: crate::protocol::TaskSupportMode,
1254    pub(crate) state: S,
1255    pub(crate) handler: F,
1256    pub(crate) _phantom: PhantomData<(T, I)>,
1257}
1258
1259#[allow(deprecated)]
1260impl<S, F, T, I> ToolBuilderWithTypedExtractor<S, F, T, I>
1261where
1262    S: Clone + Send + Sync + 'static,
1263    F: TypedExtractorHandler<S, T, I> + Clone,
1264    T: Send + Sync + 'static,
1265    I: JsonSchema + Send + Sync + 'static,
1266{
1267    /// Build the tool.
1268    pub fn build(self) -> Tool {
1269        let input_schema = {
1270            let schema = self.input_schema_override.unwrap_or_else(|| {
1271                let schema = schemars::schema_for!(I);
1272                serde_json::to_value(schema).unwrap_or_else(|_| {
1273                    serde_json::json!({
1274                        "type": "object"
1275                    })
1276                })
1277            });
1278            crate::tool::ensure_object_schema(schema)
1279        };
1280
1281        let handler = TypedExtractorToolHandler {
1282            state: self.state,
1283            handler: self.handler,
1284            input_schema: input_schema.clone(),
1285            _phantom: PhantomData,
1286        };
1287
1288        let handler_service = crate::tool::ToolHandlerService::new(handler);
1289        let catch_error = ToolCatchError::new(handler_service);
1290        let service = BoxCloneService::new(catch_error);
1291
1292        Tool {
1293            name: self.name,
1294            title: self.title,
1295            description: self.description,
1296            output_schema: self.output_schema,
1297            icons: self.icons,
1298            annotations: self.annotations,
1299            meta: None,
1300            task_support: self.task_support,
1301            required_client_capabilities: None,
1302            service: Some(service),
1303            #[cfg(feature = "stateless")]
1304            mrtr_handler: None,
1305            input_schema,
1306        }
1307    }
1308}
1309
1310/// Internal handler wrapper for typed extractor-based handlers
1311struct TypedExtractorToolHandler<S, F, T, I> {
1312    state: S,
1313    handler: F,
1314    input_schema: Value,
1315    _phantom: PhantomData<(T, I)>,
1316}
1317
1318#[allow(deprecated)]
1319impl<S, F, T, I> ToolHandler for TypedExtractorToolHandler<S, F, T, I>
1320where
1321    S: Clone + Send + Sync + 'static,
1322    F: TypedExtractorHandler<S, T, I> + Clone,
1323    T: Send + Sync + 'static,
1324    I: JsonSchema + Send + Sync + 'static,
1325{
1326    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
1327        let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
1328        self.call_with_context(ctx, args)
1329    }
1330
1331    fn call_with_context(
1332        &self,
1333        ctx: RequestContext,
1334        args: Value,
1335    ) -> BoxFuture<'_, Result<CallToolResult>> {
1336        let state = self.state.clone();
1337        let handler = self.handler.clone();
1338        Box::pin(async move { handler.call(ctx, state, args).await })
1339    }
1340
1341    fn uses_context(&self) -> bool {
1342        true
1343    }
1344
1345    fn input_schema(&self) -> Value {
1346        self.input_schema.clone()
1347    }
1348}
1349
1350#[cfg(test)]
1351mod tests {
1352    use super::*;
1353    use crate::protocol::RequestId;
1354    use schemars::JsonSchema;
1355    use serde::Deserialize;
1356    use std::sync::Arc;
1357
1358    #[derive(Debug, Deserialize, JsonSchema)]
1359    struct TestInput {
1360        name: String,
1361        count: i32,
1362    }
1363
1364    // Regression guard for the `tower_mcp::schemars` re-export (see #936).
1365    // Deriving with `#[schemars(crate = "crate::schemars")]` forces the derive
1366    // to resolve through the re-export, which is the path downstream users rely
1367    // on to stay version-aligned. `HasSchema::schema()` must then succeed.
1368    #[derive(Debug, Deserialize, JsonSchema)]
1369    #[schemars(crate = "crate::schemars")]
1370    struct ReexportInput {
1371        field: String,
1372    }
1373
1374    #[test]
1375    fn reexported_schemars_derive_produces_schema() {
1376        let schema = <Json<ReexportInput> as HasSchema>::schema()
1377            .expect("re-exported schemars derive should yield a schema");
1378        assert_eq!(schema["type"], "object");
1379        assert!(schema["properties"].get("field").is_some());
1380
1381        // Exercise the full extract path so the derived type is actually used.
1382        let ctx = RequestContext::new(RequestId::Number(1));
1383        let args = serde_json::json!({"field": "value"});
1384        let Json(input) = Json::<ReexportInput>::from_tool_request(&ctx, &(), &args)
1385            .expect("deserialization should succeed");
1386        assert_eq!(input.field, "value");
1387    }
1388
1389    #[test]
1390    fn test_json_extraction() {
1391        let args = serde_json::json!({"name": "test", "count": 42});
1392        let ctx = RequestContext::new(RequestId::Number(1));
1393
1394        let result = Json::<TestInput>::from_tool_request(&ctx, &(), &args);
1395        assert!(result.is_ok());
1396        let Json(input) = result.unwrap();
1397        assert_eq!(input.name, "test");
1398        assert_eq!(input.count, 42);
1399    }
1400
1401    #[test]
1402    fn test_json_extraction_error() {
1403        let args = serde_json::json!({"name": "test"}); // missing count
1404        let ctx = RequestContext::new(RequestId::Number(1));
1405
1406        let result = Json::<TestInput>::from_tool_request(&ctx, &(), &args);
1407        assert!(result.is_err());
1408        let rejection = result.unwrap_err();
1409        // JsonRejection contains the serde error message
1410        assert!(rejection.message().contains("count"));
1411    }
1412
1413    #[test]
1414    fn test_state_extraction() {
1415        let args = serde_json::json!({});
1416        let ctx = RequestContext::new(RequestId::Number(1));
1417        let state = Arc::new("my-state".to_string());
1418
1419        let result = State::<Arc<String>>::from_tool_request(&ctx, &state, &args);
1420        assert!(result.is_ok());
1421        let State(extracted) = result.unwrap();
1422        assert_eq!(*extracted, "my-state");
1423    }
1424
1425    #[test]
1426    fn test_context_extraction() {
1427        let args = serde_json::json!({});
1428        let ctx = RequestContext::new(RequestId::Number(42));
1429
1430        let result = Context::from_tool_request(&ctx, &(), &args);
1431        assert!(result.is_ok());
1432        let extracted = result.unwrap();
1433        assert_eq!(*extracted.request_id(), RequestId::Number(42));
1434    }
1435
1436    #[test]
1437    fn test_raw_args_extraction() {
1438        let args = serde_json::json!({"foo": "bar", "baz": 123});
1439        let ctx = RequestContext::new(RequestId::Number(1));
1440
1441        let result = RawArgs::from_tool_request(&ctx, &(), &args);
1442        assert!(result.is_ok());
1443        let RawArgs(extracted) = result.unwrap();
1444        assert_eq!(extracted["foo"], "bar");
1445        assert_eq!(extracted["baz"], 123);
1446    }
1447
1448    #[test]
1449    fn test_extension_extraction() {
1450        use crate::context::Extensions;
1451
1452        #[derive(Clone, Debug, PartialEq)]
1453        struct DatabasePool {
1454            url: String,
1455        }
1456
1457        let args = serde_json::json!({});
1458
1459        // Create extensions with a value
1460        let mut extensions = Extensions::new();
1461        extensions.insert(Arc::new(DatabasePool {
1462            url: "postgres://localhost".to_string(),
1463        }));
1464
1465        // Create context with extensions
1466        let ctx = RequestContext::new(RequestId::Number(1)).with_extensions(Arc::new(extensions));
1467
1468        // Extract the extension
1469        let result = Extension::<Arc<DatabasePool>>::from_tool_request(&ctx, &(), &args);
1470        assert!(result.is_ok());
1471        let Extension(pool) = result.unwrap();
1472        assert_eq!(pool.url, "postgres://localhost");
1473    }
1474
1475    #[test]
1476    fn test_extension_extraction_missing() {
1477        #[derive(Clone, Debug)]
1478        struct NotPresent;
1479
1480        let args = serde_json::json!({});
1481        let ctx = RequestContext::new(RequestId::Number(1));
1482
1483        // Try to extract something that's not in extensions
1484        let result = Extension::<NotPresent>::from_tool_request(&ctx, &(), &args);
1485        assert!(result.is_err());
1486        let rejection = result.unwrap_err();
1487        // ExtensionRejection contains the type name
1488        assert!(rejection.type_name().contains("NotPresent"));
1489    }
1490
1491    #[tokio::test]
1492    async fn test_single_extractor_handler() {
1493        let handler = |Json(input): Json<TestInput>| async move {
1494            Ok(CallToolResult::text(format!(
1495                "{}: {}",
1496                input.name, input.count
1497            )))
1498        };
1499
1500        let ctx = RequestContext::new(RequestId::Number(1));
1501        let args = serde_json::json!({"name": "test", "count": 5});
1502
1503        // Use explicit trait to avoid ambiguity
1504        let result: Result<CallToolResult> =
1505            ExtractorHandler::<(), (Json<TestInput>,)>::call(handler, ctx, (), args).await;
1506        assert!(result.is_ok());
1507    }
1508
1509    #[tokio::test]
1510    async fn test_two_extractor_handler() {
1511        let handler = |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1512            Ok(CallToolResult::text(format!(
1513                "{}: {} - {}",
1514                state, input.name, input.count
1515            )))
1516        };
1517
1518        let ctx = RequestContext::new(RequestId::Number(1));
1519        let state = Arc::new("prefix".to_string());
1520        let args = serde_json::json!({"name": "test", "count": 5});
1521
1522        // Use explicit trait to avoid ambiguity
1523        let result: Result<CallToolResult> = ExtractorHandler::<
1524            Arc<String>,
1525            (State<Arc<String>>, Json<TestInput>),
1526        >::call(handler, ctx, state, args)
1527        .await;
1528        assert!(result.is_ok());
1529    }
1530
1531    #[tokio::test]
1532    async fn test_three_extractor_handler() {
1533        let handler = |State(state): State<Arc<String>>,
1534                       ctx: Context,
1535                       Json(input): Json<TestInput>| async move {
1536            // Verify we can access all extractors
1537            assert!(!ctx.is_cancelled());
1538            Ok(CallToolResult::text(format!(
1539                "{}: {} - {}",
1540                state, input.name, input.count
1541            )))
1542        };
1543
1544        let ctx = RequestContext::new(RequestId::Number(1));
1545        let state = Arc::new("prefix".to_string());
1546        let args = serde_json::json!({"name": "test", "count": 5});
1547
1548        // Use explicit trait to avoid ambiguity
1549        let result: Result<CallToolResult> = ExtractorHandler::<
1550            Arc<String>,
1551            (State<Arc<String>>, Context, Json<TestInput>),
1552        >::call(handler, ctx, state, args)
1553        .await;
1554        assert!(result.is_ok());
1555    }
1556
1557    #[test]
1558    fn test_json_schema_generation() {
1559        let schema = Json::<TestInput>::schema();
1560        assert!(schema.is_some());
1561        let schema = schema.unwrap();
1562        assert!(schema.get("properties").is_some());
1563    }
1564
1565    #[test]
1566    fn test_rejection_into_error() {
1567        let rejection = Rejection::new("test error");
1568        let error: Error = rejection.into();
1569        assert!(error.to_string().contains("test error"));
1570    }
1571
1572    #[test]
1573    fn test_json_rejection() {
1574        // Test basic JsonRejection
1575        let rejection = JsonRejection::new("missing field `name`");
1576        assert_eq!(rejection.message(), "missing field `name`");
1577        assert!(rejection.path().is_none());
1578        assert!(rejection.to_string().contains("Invalid input"));
1579
1580        // Test JsonRejection with path
1581        let rejection = JsonRejection::with_path("expected string", "users[0].name");
1582        assert_eq!(rejection.message(), "expected string");
1583        assert_eq!(rejection.path(), Some("users[0].name"));
1584        assert!(rejection.to_string().contains("users[0].name"));
1585
1586        // Test conversion to Error
1587        let error: Error = rejection.into();
1588        assert!(error.to_string().contains("users[0].name"));
1589    }
1590
1591    #[test]
1592    fn test_json_rejection_from_serde_error() {
1593        // Create a real serde error by deserializing invalid JSON
1594        #[derive(Debug, serde::Deserialize)]
1595        struct TestStruct {
1596            #[allow(dead_code)]
1597            name: String,
1598        }
1599
1600        let result: std::result::Result<TestStruct, _> =
1601            serde_json::from_value(serde_json::json!({"count": 42}));
1602        assert!(result.is_err());
1603
1604        let rejection: JsonRejection = result.unwrap_err().into();
1605        assert!(rejection.message().contains("name"));
1606    }
1607
1608    #[test]
1609    fn test_extension_rejection() {
1610        // Test ExtensionRejection
1611        let rejection = ExtensionRejection::not_found::<String>();
1612        assert!(rejection.type_name().contains("String"));
1613        assert!(rejection.to_string().contains("not found"));
1614        assert!(rejection.to_string().contains("with_state"));
1615
1616        // Test conversion to Error
1617        let error: Error = rejection.into();
1618        assert!(error.to_string().contains("not found"));
1619    }
1620
1621    #[tokio::test]
1622    async fn test_tool_builder_extractor_handler() {
1623        use crate::ToolBuilder;
1624
1625        let state = Arc::new("shared-state".to_string());
1626
1627        let tool =
1628            ToolBuilder::new("test_extractor")
1629                .description("Test extractor handler")
1630                .extractor_handler(
1631                    state,
1632                    |State(state): State<Arc<String>>,
1633                     ctx: Context,
1634                     Json(input): Json<TestInput>| async move {
1635                        assert!(!ctx.is_cancelled());
1636                        Ok(CallToolResult::text(format!(
1637                            "{}: {} - {}",
1638                            state, input.name, input.count
1639                        )))
1640                    },
1641                )
1642                .build();
1643
1644        assert_eq!(tool.name, "test_extractor");
1645        assert_eq!(tool.description.as_deref(), Some("Test extractor handler"));
1646
1647        // Test calling the tool
1648        let result = tool
1649            .call(serde_json::json!({"name": "test", "count": 42}))
1650            .await;
1651        assert!(!result.is_error);
1652    }
1653
1654    #[tokio::test]
1655    #[allow(deprecated)]
1656    async fn test_tool_builder_extractor_handler_typed() {
1657        use crate::ToolBuilder;
1658
1659        let state = Arc::new("typed-state".to_string());
1660
1661        let tool = ToolBuilder::new("test_typed")
1662            .description("Test typed extractor handler")
1663            .extractor_handler_typed::<_, _, _, TestInput>(
1664                state,
1665                |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1666                    Ok(CallToolResult::text(format!(
1667                        "{}: {} - {}",
1668                        state, input.name, input.count
1669                    )))
1670                },
1671            )
1672            .build();
1673
1674        assert_eq!(tool.name, "test_typed");
1675
1676        // Verify schema is properly generated from TestInput
1677        let def = tool.definition();
1678        let schema = def.input_schema;
1679        assert!(schema.get("properties").is_some());
1680
1681        // Test calling the tool
1682        let result = tool
1683            .call(serde_json::json!({"name": "world", "count": 99}))
1684            .await;
1685        assert!(!result.is_error);
1686    }
1687
1688    #[tokio::test]
1689    async fn test_extractor_handler_auto_schema() {
1690        use crate::ToolBuilder;
1691
1692        let state = Arc::new("auto-schema".to_string());
1693
1694        // extractor_handler (not _typed) should auto-detect Json<TestInput> schema
1695        let tool = ToolBuilder::new("test_auto_schema")
1696            .description("Test auto schema detection")
1697            .extractor_handler(
1698                state,
1699                |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1700                    Ok(CallToolResult::text(format!(
1701                        "{}: {} - {}",
1702                        state, input.name, input.count
1703                    )))
1704                },
1705            )
1706            .build();
1707
1708        // Verify schema is properly generated from TestInput (not generic object)
1709        let def = tool.definition();
1710        let schema = def.input_schema;
1711        assert!(
1712            schema.get("properties").is_some(),
1713            "Schema should have properties from TestInput, got: {}",
1714            schema
1715        );
1716        let props = schema.get("properties").unwrap();
1717        assert!(
1718            props.get("name").is_some(),
1719            "Schema should have 'name' property"
1720        );
1721        assert!(
1722            props.get("count").is_some(),
1723            "Schema should have 'count' property"
1724        );
1725
1726        // Test calling the tool
1727        let result = tool
1728            .call(serde_json::json!({"name": "world", "count": 99}))
1729            .await;
1730        assert!(!result.is_error);
1731    }
1732
1733    #[test]
1734    fn test_extractor_handler_no_json_fallback() {
1735        use crate::ToolBuilder;
1736
1737        // extractor_handler without Json<T> should fall back to generic schema
1738        let tool = ToolBuilder::new("test_no_json")
1739            .description("Test no json fallback")
1740            .extractor_handler((), |RawArgs(args): RawArgs| async move {
1741                Ok(CallToolResult::json(args))
1742            })
1743            .build();
1744
1745        let def = tool.definition();
1746        let schema = def.input_schema;
1747        assert_eq!(
1748            schema.get("type").and_then(|v| v.as_str()),
1749            Some("object"),
1750            "Schema should be generic object"
1751        );
1752        assert_eq!(
1753            schema.get("additionalProperties").and_then(|v| v.as_bool()),
1754            Some(true),
1755            "Schema should allow additional properties"
1756        );
1757        // Should NOT have specific properties
1758        assert!(
1759            schema.get("properties").is_none(),
1760            "Generic schema should not have specific properties"
1761        );
1762    }
1763
1764    #[tokio::test]
1765    async fn test_extractor_handler_with_layer() {
1766        use crate::ToolBuilder;
1767        use std::time::Duration;
1768        use tower::timeout::TimeoutLayer;
1769
1770        let state = Arc::new("layered".to_string());
1771
1772        let tool = ToolBuilder::new("test_extractor_layer")
1773            .description("Test extractor handler with layer")
1774            .extractor_handler(
1775                state,
1776                |State(s): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1777                    Ok(CallToolResult::text(format!("{}: {}", s, input.name)))
1778                },
1779            )
1780            .layer(TimeoutLayer::new(Duration::from_secs(5)))
1781            .build();
1782
1783        // Verify the tool works
1784        let result = tool
1785            .call(serde_json::json!({"name": "test", "count": 1}))
1786            .await;
1787        assert!(!result.is_error);
1788        assert_eq!(result.first_text().unwrap(), "layered: test");
1789
1790        // Verify schema is still properly generated
1791        let def = tool.definition();
1792        let schema = def.input_schema;
1793        assert!(
1794            schema.get("properties").is_some(),
1795            "Schema should have properties even with layer"
1796        );
1797    }
1798
1799    #[tokio::test]
1800    async fn test_extractor_handler_with_timeout_layer() {
1801        use crate::ToolBuilder;
1802        use std::time::Duration;
1803        use tower::timeout::TimeoutLayer;
1804
1805        let tool = ToolBuilder::new("test_extractor_timeout")
1806            .description("Test extractor handler timeout")
1807            .extractor_handler((), |Json(input): Json<TestInput>| async move {
1808                tokio::time::sleep(Duration::from_millis(200)).await;
1809                Ok(CallToolResult::text(input.name.to_string()))
1810            })
1811            .layer(TimeoutLayer::new(Duration::from_millis(50)))
1812            .build();
1813
1814        // Should timeout
1815        let result = tool
1816            .call(serde_json::json!({"name": "slow", "count": 1}))
1817            .await;
1818        assert!(result.is_error);
1819        let msg = result.first_text().unwrap().to_lowercase();
1820        assert!(
1821            msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
1822            "Expected timeout error, got: {}",
1823            msg
1824        );
1825    }
1826
1827    #[tokio::test]
1828    async fn test_extractor_handler_with_multiple_layers() {
1829        use crate::ToolBuilder;
1830        use std::time::Duration;
1831        use tower::limit::ConcurrencyLimitLayer;
1832        use tower::timeout::TimeoutLayer;
1833
1834        let state = Arc::new("multi".to_string());
1835
1836        let tool = ToolBuilder::new("test_multi_layer")
1837            .description("Test multiple layers")
1838            .extractor_handler(
1839                state,
1840                |State(s): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1841                    Ok(CallToolResult::text(format!("{}: {}", s, input.name)))
1842                },
1843            )
1844            .layer(TimeoutLayer::new(Duration::from_secs(5)))
1845            .layer(ConcurrencyLimitLayer::new(10))
1846            .build();
1847
1848        let result = tool
1849            .call(serde_json::json!({"name": "test", "count": 1}))
1850            .await;
1851        assert!(!result.is_error);
1852        assert_eq!(result.first_text().unwrap(), "multi: test");
1853    }
1854}
tower_mcp/extract.rs

tower_mcp/
extract.rs
