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

tower_mcp/
extract.rs
