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

tower_mcp/
extract.rs
