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            task_preparer: None,
1067            service: Some(service),
1068            #[cfg(feature = "stateless")]
1069            mrtr_handler: None,
1070            input_schema: self.input_schema,
1071        }
1072    }
1073
1074    /// Apply a Tower layer (middleware) to this tool.
1075    ///
1076    /// The layer wraps the tool's handler service, enabling functionality like
1077    /// timeouts, rate limiting, and metrics collection at the per-tool level.
1078    ///
1079    /// # Example
1080    ///
1081    /// ```rust
1082    /// use std::sync::Arc;
1083    /// use std::time::Duration;
1084    /// use tower::timeout::TimeoutLayer;
1085    /// use tower_mcp::{ToolBuilder, CallToolResult};
1086    /// use tower_mcp::extract::{Json, State};
1087    /// use schemars::JsonSchema;
1088    /// use serde::Deserialize;
1089    ///
1090    /// #[derive(Clone)]
1091    /// struct AppState { prefix: String }
1092    ///
1093    /// #[derive(Debug, Deserialize, JsonSchema)]
1094    /// struct QueryInput { query: String }
1095    ///
1096    /// let state = Arc::new(AppState { prefix: "db".to_string() });
1097    ///
1098    /// let tool = ToolBuilder::new("search")
1099    ///     .description("Search with timeout")
1100    ///     .extractor_handler(state, |
1101    ///         State(app): State<Arc<AppState>>,
1102    ///         Json(input): Json<QueryInput>,
1103    ///     | async move {
1104    ///         Ok(CallToolResult::text(format!("{}: {}", app.prefix, input.query)))
1105    ///     })
1106    ///     .layer(TimeoutLayer::new(Duration::from_secs(30)))
1107    ///     .build();
1108    /// ```
1109    pub fn layer<L>(self, layer: L) -> ToolBuilderWithExtractorLayer<S, F, T, L> {
1110        ToolBuilderWithExtractorLayer {
1111            name: self.name,
1112            title: self.title,
1113            description: self.description,
1114            output_schema: self.output_schema,
1115            icons: self.icons,
1116            annotations: self.annotations,
1117            task_support: self.task_support,
1118            state: self.state,
1119            handler: self.handler,
1120            input_schema: self.input_schema,
1121            layer,
1122            _phantom: PhantomData,
1123        }
1124    }
1125
1126    /// Apply a guard to this tool.
1127    ///
1128    /// See [`ToolBuilderWithHandler::guard`](crate::ToolBuilder) for details.
1129    pub fn guard<G>(self, guard: G) -> ToolBuilderWithExtractorLayer<S, F, T, GuardLayer<G>>
1130    where
1131        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1132    {
1133        self.layer(GuardLayer::new(guard))
1134    }
1135}
1136
1137/// Builder state after a layer has been applied to an extractor handler.
1138///
1139/// This builder allows chaining additional layers and building the final tool.
1140#[doc(hidden)]
1141pub struct ToolBuilderWithExtractorLayer<S, F, T, L> {
1142    name: String,
1143    title: Option<String>,
1144    description: Option<String>,
1145    output_schema: Option<Value>,
1146    icons: Option<Vec<crate::protocol::ToolIcon>>,
1147    annotations: Option<crate::protocol::ToolAnnotations>,
1148    task_support: crate::protocol::TaskSupportMode,
1149    state: S,
1150    handler: F,
1151    input_schema: Value,
1152    layer: L,
1153    _phantom: PhantomData<T>,
1154}
1155
1156#[allow(private_bounds)]
1157impl<S, F, T, L> ToolBuilderWithExtractorLayer<S, F, T, L>
1158where
1159    S: Clone + Send + Sync + 'static,
1160    F: ExtractorHandler<S, T> + Clone,
1161    T: Send + Sync + 'static,
1162    L: tower::Layer<ToolHandlerService<ExtractorToolHandler<S, F, T>>>
1163        + Clone
1164        + Send
1165        + Sync
1166        + 'static,
1167    L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
1168    <L::Service as Service<ToolRequest>>::Error: std::fmt::Display + Send,
1169    <L::Service as Service<ToolRequest>>::Future: Send,
1170{
1171    /// Build the tool with the applied layer(s).
1172    pub fn build(self) -> Tool {
1173        let handler = ExtractorToolHandler {
1174            state: self.state,
1175            handler: self.handler,
1176            input_schema: self.input_schema.clone(),
1177            _phantom: PhantomData,
1178        };
1179
1180        let handler_service = ToolHandlerService::new(handler);
1181        let layered = self.layer.layer(handler_service);
1182        let catch_error = ToolCatchError::new(layered);
1183        let service = BoxCloneService::new(catch_error);
1184
1185        Tool {
1186            name: self.name,
1187            title: self.title,
1188            description: self.description,
1189            output_schema: self.output_schema,
1190            icons: self.icons,
1191            annotations: self.annotations,
1192            meta: None,
1193            task_support: self.task_support,
1194            required_client_capabilities: None,
1195            task_preparer: None,
1196            service: Some(service),
1197            #[cfg(feature = "stateless")]
1198            mrtr_handler: None,
1199            input_schema: self.input_schema,
1200        }
1201    }
1202
1203    /// Apply an additional Tower layer (middleware).
1204    ///
1205    /// Layers are applied in order, with earlier layers wrapping later ones.
1206    /// This means the first layer added is the outermost middleware.
1207    pub fn layer<L2>(
1208        self,
1209        layer: L2,
1210    ) -> ToolBuilderWithExtractorLayer<S, F, T, tower::layer::util::Stack<L2, L>> {
1211        ToolBuilderWithExtractorLayer {
1212            name: self.name,
1213            title: self.title,
1214            description: self.description,
1215            output_schema: self.output_schema,
1216            icons: self.icons,
1217            annotations: self.annotations,
1218            task_support: self.task_support,
1219            state: self.state,
1220            handler: self.handler,
1221            input_schema: self.input_schema,
1222            layer: tower::layer::util::Stack::new(layer, self.layer),
1223            _phantom: PhantomData,
1224        }
1225    }
1226
1227    /// Apply a guard to this tool.
1228    ///
1229    /// See [`ToolBuilderWithHandler::guard`](crate::ToolBuilder) for details.
1230    pub fn guard<G>(
1231        self,
1232        guard: G,
1233    ) -> ToolBuilderWithExtractorLayer<S, F, T, tower::layer::util::Stack<GuardLayer<G>, L>>
1234    where
1235        G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1236    {
1237        self.layer(GuardLayer::new(guard))
1238    }
1239}
1240
1241/// Builder state for extractor-based handlers with typed JSON input
1242#[doc(hidden)]
1243#[deprecated(
1244    since = "0.8.0",
1245    note = "Use `ToolBuilderWithExtractor` via `extractor_handler` instead"
1246)]
1247pub struct ToolBuilderWithTypedExtractor<S, F, T, I> {
1248    pub(crate) name: String,
1249    pub(crate) title: Option<String>,
1250    pub(crate) description: Option<String>,
1251    pub(crate) output_schema: Option<Value>,
1252    pub(crate) input_schema_override: Option<Value>,
1253    pub(crate) icons: Option<Vec<crate::protocol::ToolIcon>>,
1254    pub(crate) annotations: Option<crate::protocol::ToolAnnotations>,
1255    pub(crate) task_support: crate::protocol::TaskSupportMode,
1256    pub(crate) state: S,
1257    pub(crate) handler: F,
1258    pub(crate) _phantom: PhantomData<(T, I)>,
1259}
1260
1261#[allow(deprecated)]
1262impl<S, F, T, I> ToolBuilderWithTypedExtractor<S, F, T, I>
1263where
1264    S: Clone + Send + Sync + 'static,
1265    F: TypedExtractorHandler<S, T, I> + Clone,
1266    T: Send + Sync + 'static,
1267    I: JsonSchema + Send + Sync + 'static,
1268{
1269    /// Build the tool.
1270    pub fn build(self) -> Tool {
1271        let input_schema = {
1272            let schema = self.input_schema_override.unwrap_or_else(|| {
1273                let schema = schemars::schema_for!(I);
1274                serde_json::to_value(schema).unwrap_or_else(|_| {
1275                    serde_json::json!({
1276                        "type": "object"
1277                    })
1278                })
1279            });
1280            crate::tool::ensure_object_schema(schema)
1281        };
1282
1283        let handler = TypedExtractorToolHandler {
1284            state: self.state,
1285            handler: self.handler,
1286            input_schema: input_schema.clone(),
1287            _phantom: PhantomData,
1288        };
1289
1290        let handler_service = crate::tool::ToolHandlerService::new(handler);
1291        let catch_error = ToolCatchError::new(handler_service);
1292        let service = BoxCloneService::new(catch_error);
1293
1294        Tool {
1295            name: self.name,
1296            title: self.title,
1297            description: self.description,
1298            output_schema: self.output_schema,
1299            icons: self.icons,
1300            annotations: self.annotations,
1301            meta: None,
1302            task_support: self.task_support,
1303            required_client_capabilities: None,
1304            task_preparer: None,
1305            service: Some(service),
1306            #[cfg(feature = "stateless")]
1307            mrtr_handler: None,
1308            input_schema,
1309        }
1310    }
1311}
1312
1313/// Internal handler wrapper for typed extractor-based handlers
1314struct TypedExtractorToolHandler<S, F, T, I> {
1315    state: S,
1316    handler: F,
1317    input_schema: Value,
1318    _phantom: PhantomData<(T, I)>,
1319}
1320
1321#[allow(deprecated)]
1322impl<S, F, T, I> ToolHandler for TypedExtractorToolHandler<S, F, T, I>
1323where
1324    S: Clone + Send + Sync + 'static,
1325    F: TypedExtractorHandler<S, T, I> + Clone,
1326    T: Send + Sync + 'static,
1327    I: JsonSchema + Send + Sync + 'static,
1328{
1329    fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
1330        let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
1331        self.call_with_context(ctx, args)
1332    }
1333
1334    fn call_with_context(
1335        &self,
1336        ctx: RequestContext,
1337        args: Value,
1338    ) -> BoxFuture<'_, Result<CallToolResult>> {
1339        let state = self.state.clone();
1340        let handler = self.handler.clone();
1341        Box::pin(async move { handler.call(ctx, state, args).await })
1342    }
1343
1344    fn uses_context(&self) -> bool {
1345        true
1346    }
1347
1348    fn input_schema(&self) -> Value {
1349        self.input_schema.clone()
1350    }
1351}
1352
1353#[cfg(test)]
1354mod tests {
1355    use super::*;
1356    use crate::protocol::RequestId;
1357    use schemars::JsonSchema;
1358    use serde::Deserialize;
1359    use std::sync::Arc;
1360
1361    #[derive(Debug, Deserialize, JsonSchema)]
1362    struct TestInput {
1363        name: String,
1364        count: i32,
1365    }
1366
1367    // Regression guard for the `tower_mcp::schemars` re-export (see #936).
1368    // Deriving with `#[schemars(crate = "crate::schemars")]` forces the derive
1369    // to resolve through the re-export, which is the path downstream users rely
1370    // on to stay version-aligned. `HasSchema::schema()` must then succeed.
1371    #[derive(Debug, Deserialize, JsonSchema)]
1372    #[schemars(crate = "crate::schemars")]
1373    struct ReexportInput {
1374        field: String,
1375    }
1376
1377    #[test]
1378    fn reexported_schemars_derive_produces_schema() {
1379        let schema = <Json<ReexportInput> as HasSchema>::schema()
1380            .expect("re-exported schemars derive should yield a schema");
1381        assert_eq!(schema["type"], "object");
1382        assert!(schema["properties"].get("field").is_some());
1383
1384        // Exercise the full extract path so the derived type is actually used.
1385        let ctx = RequestContext::new(RequestId::Number(1));
1386        let args = serde_json::json!({"field": "value"});
1387        let Json(input) = Json::<ReexportInput>::from_tool_request(&ctx, &(), &args)
1388            .expect("deserialization should succeed");
1389        assert_eq!(input.field, "value");
1390    }
1391
1392    #[test]
1393    fn test_json_extraction() {
1394        let args = serde_json::json!({"name": "test", "count": 42});
1395        let ctx = RequestContext::new(RequestId::Number(1));
1396
1397        let result = Json::<TestInput>::from_tool_request(&ctx, &(), &args);
1398        assert!(result.is_ok());
1399        let Json(input) = result.unwrap();
1400        assert_eq!(input.name, "test");
1401        assert_eq!(input.count, 42);
1402    }
1403
1404    #[test]
1405    fn test_json_extraction_error() {
1406        let args = serde_json::json!({"name": "test"}); // missing count
1407        let ctx = RequestContext::new(RequestId::Number(1));
1408
1409        let result = Json::<TestInput>::from_tool_request(&ctx, &(), &args);
1410        assert!(result.is_err());
1411        let rejection = result.unwrap_err();
1412        // JsonRejection contains the serde error message
1413        assert!(rejection.message().contains("count"));
1414    }
1415
1416    #[test]
1417    fn test_state_extraction() {
1418        let args = serde_json::json!({});
1419        let ctx = RequestContext::new(RequestId::Number(1));
1420        let state = Arc::new("my-state".to_string());
1421
1422        let result = State::<Arc<String>>::from_tool_request(&ctx, &state, &args);
1423        assert!(result.is_ok());
1424        let State(extracted) = result.unwrap();
1425        assert_eq!(*extracted, "my-state");
1426    }
1427
1428    #[test]
1429    fn test_context_extraction() {
1430        let args = serde_json::json!({});
1431        let ctx = RequestContext::new(RequestId::Number(42));
1432
1433        let result = Context::from_tool_request(&ctx, &(), &args);
1434        assert!(result.is_ok());
1435        let extracted = result.unwrap();
1436        assert_eq!(*extracted.request_id(), RequestId::Number(42));
1437    }
1438
1439    #[test]
1440    fn test_raw_args_extraction() {
1441        let args = serde_json::json!({"foo": "bar", "baz": 123});
1442        let ctx = RequestContext::new(RequestId::Number(1));
1443
1444        let result = RawArgs::from_tool_request(&ctx, &(), &args);
1445        assert!(result.is_ok());
1446        let RawArgs(extracted) = result.unwrap();
1447        assert_eq!(extracted["foo"], "bar");
1448        assert_eq!(extracted["baz"], 123);
1449    }
1450
1451    #[test]
1452    fn test_extension_extraction() {
1453        use crate::context::Extensions;
1454
1455        #[derive(Clone, Debug, PartialEq)]
1456        struct DatabasePool {
1457            url: String,
1458        }
1459
1460        let args = serde_json::json!({});
1461
1462        // Create extensions with a value
1463        let mut extensions = Extensions::new();
1464        extensions.insert(Arc::new(DatabasePool {
1465            url: "postgres://localhost".to_string(),
1466        }));
1467
1468        // Create context with extensions
1469        let ctx = RequestContext::new(RequestId::Number(1)).with_extensions(Arc::new(extensions));
1470
1471        // Extract the extension
1472        let result = Extension::<Arc<DatabasePool>>::from_tool_request(&ctx, &(), &args);
1473        assert!(result.is_ok());
1474        let Extension(pool) = result.unwrap();
1475        assert_eq!(pool.url, "postgres://localhost");
1476    }
1477
1478    #[test]
1479    fn test_extension_extraction_missing() {
1480        #[derive(Clone, Debug)]
1481        struct NotPresent;
1482
1483        let args = serde_json::json!({});
1484        let ctx = RequestContext::new(RequestId::Number(1));
1485
1486        // Try to extract something that's not in extensions
1487        let result = Extension::<NotPresent>::from_tool_request(&ctx, &(), &args);
1488        assert!(result.is_err());
1489        let rejection = result.unwrap_err();
1490        // ExtensionRejection contains the type name
1491        assert!(rejection.type_name().contains("NotPresent"));
1492    }
1493
1494    #[tokio::test]
1495    async fn test_single_extractor_handler() {
1496        let handler = |Json(input): Json<TestInput>| async move {
1497            Ok(CallToolResult::text(format!(
1498                "{}: {}",
1499                input.name, input.count
1500            )))
1501        };
1502
1503        let ctx = RequestContext::new(RequestId::Number(1));
1504        let args = serde_json::json!({"name": "test", "count": 5});
1505
1506        // Use explicit trait to avoid ambiguity
1507        let result: Result<CallToolResult> =
1508            ExtractorHandler::<(), (Json<TestInput>,)>::call(handler, ctx, (), args).await;
1509        assert!(result.is_ok());
1510    }
1511
1512    #[tokio::test]
1513    async fn test_two_extractor_handler() {
1514        let handler = |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1515            Ok(CallToolResult::text(format!(
1516                "{}: {} - {}",
1517                state, input.name, input.count
1518            )))
1519        };
1520
1521        let ctx = RequestContext::new(RequestId::Number(1));
1522        let state = Arc::new("prefix".to_string());
1523        let args = serde_json::json!({"name": "test", "count": 5});
1524
1525        // Use explicit trait to avoid ambiguity
1526        let result: Result<CallToolResult> = ExtractorHandler::<
1527            Arc<String>,
1528            (State<Arc<String>>, Json<TestInput>),
1529        >::call(handler, ctx, state, args)
1530        .await;
1531        assert!(result.is_ok());
1532    }
1533
1534    #[tokio::test]
1535    async fn test_three_extractor_handler() {
1536        let handler = |State(state): State<Arc<String>>,
1537                       ctx: Context,
1538                       Json(input): Json<TestInput>| async move {
1539            // Verify we can access all extractors
1540            assert!(!ctx.is_cancelled());
1541            Ok(CallToolResult::text(format!(
1542                "{}: {} - {}",
1543                state, input.name, input.count
1544            )))
1545        };
1546
1547        let ctx = RequestContext::new(RequestId::Number(1));
1548        let state = Arc::new("prefix".to_string());
1549        let args = serde_json::json!({"name": "test", "count": 5});
1550
1551        // Use explicit trait to avoid ambiguity
1552        let result: Result<CallToolResult> = ExtractorHandler::<
1553            Arc<String>,
1554            (State<Arc<String>>, Context, Json<TestInput>),
1555        >::call(handler, ctx, state, args)
1556        .await;
1557        assert!(result.is_ok());
1558    }
1559
1560    #[test]
1561    fn test_json_schema_generation() {
1562        let schema = Json::<TestInput>::schema();
1563        assert!(schema.is_some());
1564        let schema = schema.unwrap();
1565        assert!(schema.get("properties").is_some());
1566    }
1567
1568    #[test]
1569    fn test_rejection_into_error() {
1570        let rejection = Rejection::new("test error");
1571        let error: Error = rejection.into();
1572        assert!(error.to_string().contains("test error"));
1573    }
1574
1575    #[test]
1576    fn test_json_rejection() {
1577        // Test basic JsonRejection
1578        let rejection = JsonRejection::new("missing field `name`");
1579        assert_eq!(rejection.message(), "missing field `name`");
1580        assert!(rejection.path().is_none());
1581        assert!(rejection.to_string().contains("Invalid input"));
1582
1583        // Test JsonRejection with path
1584        let rejection = JsonRejection::with_path("expected string", "users[0].name");
1585        assert_eq!(rejection.message(), "expected string");
1586        assert_eq!(rejection.path(), Some("users[0].name"));
1587        assert!(rejection.to_string().contains("users[0].name"));
1588
1589        // Test conversion to Error
1590        let error: Error = rejection.into();
1591        assert!(error.to_string().contains("users[0].name"));
1592    }
1593
1594    #[test]
1595    fn test_json_rejection_from_serde_error() {
1596        // Create a real serde error by deserializing invalid JSON
1597        #[derive(Debug, serde::Deserialize)]
1598        struct TestStruct {
1599            #[allow(dead_code)]
1600            name: String,
1601        }
1602
1603        let result: std::result::Result<TestStruct, _> =
1604            serde_json::from_value(serde_json::json!({"count": 42}));
1605        assert!(result.is_err());
1606
1607        let rejection: JsonRejection = result.unwrap_err().into();
1608        assert!(rejection.message().contains("name"));
1609    }
1610
1611    #[test]
1612    fn test_extension_rejection() {
1613        // Test ExtensionRejection
1614        let rejection = ExtensionRejection::not_found::<String>();
1615        assert!(rejection.type_name().contains("String"));
1616        assert!(rejection.to_string().contains("not found"));
1617        assert!(rejection.to_string().contains("with_state"));
1618
1619        // Test conversion to Error
1620        let error: Error = rejection.into();
1621        assert!(error.to_string().contains("not found"));
1622    }
1623
1624    #[tokio::test]
1625    async fn test_tool_builder_extractor_handler() {
1626        use crate::ToolBuilder;
1627
1628        let state = Arc::new("shared-state".to_string());
1629
1630        let tool =
1631            ToolBuilder::new("test_extractor")
1632                .description("Test extractor handler")
1633                .extractor_handler(
1634                    state,
1635                    |State(state): State<Arc<String>>,
1636                     ctx: Context,
1637                     Json(input): Json<TestInput>| async move {
1638                        assert!(!ctx.is_cancelled());
1639                        Ok(CallToolResult::text(format!(
1640                            "{}: {} - {}",
1641                            state, input.name, input.count
1642                        )))
1643                    },
1644                )
1645                .build();
1646
1647        assert_eq!(tool.name, "test_extractor");
1648        assert_eq!(tool.description.as_deref(), Some("Test extractor handler"));
1649
1650        // Test calling the tool
1651        let result = tool
1652            .call(serde_json::json!({"name": "test", "count": 42}))
1653            .await;
1654        assert!(!result.is_error);
1655    }
1656
1657    #[tokio::test]
1658    #[allow(deprecated)]
1659    async fn test_tool_builder_extractor_handler_typed() {
1660        use crate::ToolBuilder;
1661
1662        let state = Arc::new("typed-state".to_string());
1663
1664        let tool = ToolBuilder::new("test_typed")
1665            .description("Test typed extractor handler")
1666            .extractor_handler_typed::<_, _, _, TestInput>(
1667                state,
1668                |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1669                    Ok(CallToolResult::text(format!(
1670                        "{}: {} - {}",
1671                        state, input.name, input.count
1672                    )))
1673                },
1674            )
1675            .build();
1676
1677        assert_eq!(tool.name, "test_typed");
1678
1679        // Verify schema is properly generated from TestInput
1680        let def = tool.definition();
1681        let schema = def.input_schema;
1682        assert!(schema.get("properties").is_some());
1683
1684        // Test calling the tool
1685        let result = tool
1686            .call(serde_json::json!({"name": "world", "count": 99}))
1687            .await;
1688        assert!(!result.is_error);
1689    }
1690
1691    #[tokio::test]
1692    async fn test_extractor_handler_auto_schema() {
1693        use crate::ToolBuilder;
1694
1695        let state = Arc::new("auto-schema".to_string());
1696
1697        // extractor_handler (not _typed) should auto-detect Json<TestInput> schema
1698        let tool = ToolBuilder::new("test_auto_schema")
1699            .description("Test auto schema detection")
1700            .extractor_handler(
1701                state,
1702                |State(state): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1703                    Ok(CallToolResult::text(format!(
1704                        "{}: {} - {}",
1705                        state, input.name, input.count
1706                    )))
1707                },
1708            )
1709            .build();
1710
1711        // Verify schema is properly generated from TestInput (not generic object)
1712        let def = tool.definition();
1713        let schema = def.input_schema;
1714        assert!(
1715            schema.get("properties").is_some(),
1716            "Schema should have properties from TestInput, got: {}",
1717            schema
1718        );
1719        let props = schema.get("properties").unwrap();
1720        assert!(
1721            props.get("name").is_some(),
1722            "Schema should have 'name' property"
1723        );
1724        assert!(
1725            props.get("count").is_some(),
1726            "Schema should have 'count' property"
1727        );
1728
1729        // Test calling the tool
1730        let result = tool
1731            .call(serde_json::json!({"name": "world", "count": 99}))
1732            .await;
1733        assert!(!result.is_error);
1734    }
1735
1736    #[test]
1737    fn test_extractor_handler_no_json_fallback() {
1738        use crate::ToolBuilder;
1739
1740        // extractor_handler without Json<T> should fall back to generic schema
1741        let tool = ToolBuilder::new("test_no_json")
1742            .description("Test no json fallback")
1743            .extractor_handler((), |RawArgs(args): RawArgs| async move {
1744                Ok(CallToolResult::json(args))
1745            })
1746            .build();
1747
1748        let def = tool.definition();
1749        let schema = def.input_schema;
1750        assert_eq!(
1751            schema.get("type").and_then(|v| v.as_str()),
1752            Some("object"),
1753            "Schema should be generic object"
1754        );
1755        assert_eq!(
1756            schema.get("additionalProperties").and_then(|v| v.as_bool()),
1757            Some(true),
1758            "Schema should allow additional properties"
1759        );
1760        // Should NOT have specific properties
1761        assert!(
1762            schema.get("properties").is_none(),
1763            "Generic schema should not have specific properties"
1764        );
1765    }
1766
1767    #[tokio::test]
1768    async fn test_extractor_handler_with_layer() {
1769        use crate::ToolBuilder;
1770        use std::time::Duration;
1771        use tower::timeout::TimeoutLayer;
1772
1773        let state = Arc::new("layered".to_string());
1774
1775        let tool = ToolBuilder::new("test_extractor_layer")
1776            .description("Test extractor handler with layer")
1777            .extractor_handler(
1778                state,
1779                |State(s): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1780                    Ok(CallToolResult::text(format!("{}: {}", s, input.name)))
1781                },
1782            )
1783            .layer(TimeoutLayer::new(Duration::from_secs(5)))
1784            .build();
1785
1786        // Verify the tool works
1787        let result = tool
1788            .call(serde_json::json!({"name": "test", "count": 1}))
1789            .await;
1790        assert!(!result.is_error);
1791        assert_eq!(result.first_text().unwrap(), "layered: test");
1792
1793        // Verify schema is still properly generated
1794        let def = tool.definition();
1795        let schema = def.input_schema;
1796        assert!(
1797            schema.get("properties").is_some(),
1798            "Schema should have properties even with layer"
1799        );
1800    }
1801
1802    #[tokio::test]
1803    async fn test_extractor_handler_with_timeout_layer() {
1804        use crate::ToolBuilder;
1805        use std::time::Duration;
1806        use tower::timeout::TimeoutLayer;
1807
1808        let tool = ToolBuilder::new("test_extractor_timeout")
1809            .description("Test extractor handler timeout")
1810            .extractor_handler((), |Json(input): Json<TestInput>| async move {
1811                tokio::time::sleep(Duration::from_millis(200)).await;
1812                Ok(CallToolResult::text(input.name.to_string()))
1813            })
1814            .layer(TimeoutLayer::new(Duration::from_millis(50)))
1815            .build();
1816
1817        // Should timeout
1818        let result = tool
1819            .call(serde_json::json!({"name": "slow", "count": 1}))
1820            .await;
1821        assert!(result.is_error);
1822        let msg = result.first_text().unwrap().to_lowercase();
1823        assert!(
1824            msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
1825            "Expected timeout error, got: {}",
1826            msg
1827        );
1828    }
1829
1830    #[tokio::test]
1831    async fn test_extractor_handler_with_multiple_layers() {
1832        use crate::ToolBuilder;
1833        use std::time::Duration;
1834        use tower::limit::ConcurrencyLimitLayer;
1835        use tower::timeout::TimeoutLayer;
1836
1837        let state = Arc::new("multi".to_string());
1838
1839        let tool = ToolBuilder::new("test_multi_layer")
1840            .description("Test multiple layers")
1841            .extractor_handler(
1842                state,
1843                |State(s): State<Arc<String>>, Json(input): Json<TestInput>| async move {
1844                    Ok(CallToolResult::text(format!("{}: {}", s, input.name)))
1845                },
1846            )
1847            .layer(TimeoutLayer::new(Duration::from_secs(5)))
1848            .layer(ConcurrencyLimitLayer::new(10))
1849            .build();
1850
1851        let result = tool
1852            .call(serde_json::json!({"name": "test", "count": 1}))
1853            .await;
1854        assert!(!result.is_error);
1855        assert_eq!(result.first_text().unwrap(), "multi: test");
1856    }
1857}
tower_mcp/extract.rs

tower_mcp/
extract.rs
