1use std::borrow::Cow;
34use std::convert::Infallible;
35use std::fmt;
36use std::future::Future;
37use std::pin::Pin;
38use std::sync::Arc;
39use std::task::{Context, Poll};
40
41use pin_project_lite::pin_project;
42
43use schemars::{JsonSchema, Schema, SchemaGenerator};
44use serde::Serialize;
45use serde::de::DeserializeOwned;
46use serde_json::Value;
47#[cfg(feature = "stateless")]
48use tower::ServiceExt;
49use tower::util::BoxCloneService;
50use tower_service::Service;
51
52#[cfg(feature = "stateless")]
53use tokio::sync::Mutex;
54
55use crate::context::RequestContext;
56use crate::error::{Error, Result, ResultExt};
57use crate::protocol::{
58 CallToolResult, ClientCapabilities, RequestOutcome, TaskSupportMode, ToolAnnotations,
59 ToolDefinition, ToolExecution, ToolIcon,
60};
61
62#[derive(Debug, Clone)]
71pub struct ToolRequest {
72 pub ctx: RequestContext,
74 pub args: Value,
76}
77
78impl ToolRequest {
79 pub fn new(ctx: RequestContext, args: Value) -> Self {
81 Self { ctx, args }
82 }
83}
84
85pub type BoxToolService = BoxCloneService<ToolRequest, CallToolResult, Infallible>;
91
92#[cfg(feature = "stateless")]
94type BoxMrtrToolService = BoxCloneService<ToolRequest, RequestOutcome<CallToolResult>, Infallible>;
95
96#[doc(hidden)]
102pub struct ToolCatchError<S> {
103 inner: S,
104}
105
106impl<S> ToolCatchError<S> {
107 pub fn new(inner: S) -> Self {
109 Self { inner }
110 }
111}
112
113impl<S: Clone> Clone for ToolCatchError<S> {
114 fn clone(&self) -> Self {
115 Self {
116 inner: self.inner.clone(),
117 }
118 }
119}
120
121impl<S: fmt::Debug> fmt::Debug for ToolCatchError<S> {
122 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123 f.debug_struct("ToolCatchError")
124 .field("inner", &self.inner)
125 .finish()
126 }
127}
128
129pin_project! {
130 #[doc(hidden)]
132 pub struct ToolCatchErrorFuture<F> {
133 #[pin]
134 inner: F,
135 }
136}
137
138impl<F, E> Future for ToolCatchErrorFuture<F>
139where
140 F: Future<Output = std::result::Result<CallToolResult, E>>,
141 E: fmt::Display,
142{
143 type Output = std::result::Result<CallToolResult, Infallible>;
144
145 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
146 match self.project().inner.poll(cx) {
147 Poll::Pending => Poll::Pending,
148 Poll::Ready(Ok(result)) => Poll::Ready(Ok(result)),
149 Poll::Ready(Err(err)) => Poll::Ready(Ok(CallToolResult::error(err.to_string()))),
150 }
151 }
152}
153
154impl<S> Service<ToolRequest> for ToolCatchError<S>
155where
156 S: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
157 S::Error: fmt::Display + Send,
158 S::Future: Send,
159{
160 type Response = CallToolResult;
161 type Error = Infallible;
162 type Future = ToolCatchErrorFuture<S::Future>;
163
164 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
165 match self.inner.poll_ready(cx) {
167 Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
168 Poll::Ready(Err(_)) => Poll::Ready(Ok(())),
169 Poll::Pending => Poll::Pending,
170 }
171 }
172
173 fn call(&mut self, req: ToolRequest) -> Self::Future {
174 ToolCatchErrorFuture {
175 inner: self.inner.call(req),
176 }
177 }
178}
179
180#[cfg(feature = "stateless")]
186#[derive(Clone)]
187struct MrtrToolCatchError<S> {
188 inner: S,
189}
190
191#[cfg(feature = "stateless")]
192impl<S> MrtrToolCatchError<S> {
193 fn new(inner: S) -> Self {
194 Self { inner }
195 }
196}
197
198#[cfg(feature = "stateless")]
199impl<S> Service<ToolRequest> for MrtrToolCatchError<S>
200where
201 S: Service<ToolRequest, Response = RequestOutcome<CallToolResult>> + Clone + Send + 'static,
202 S::Error: fmt::Display + Send + 'static,
203 S::Future: Send + 'static,
204{
205 type Response = RequestOutcome<CallToolResult>;
206 type Error = Infallible;
207 type Future =
208 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
209
210 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
211 match self.inner.poll_ready(cx) {
212 Poll::Ready(Ok(())) | Poll::Ready(Err(_)) => Poll::Ready(Ok(())),
213 Poll::Pending => Poll::Pending,
214 }
215 }
216
217 fn call(&mut self, req: ToolRequest) -> Self::Future {
218 let future = self.inner.call(req);
219 Box::pin(async move {
220 Ok(match future.await {
221 Ok(outcome) => outcome,
222 Err(error) => RequestOutcome::Complete(CallToolResult::error(error.to_string())),
223 })
224 })
225 }
226}
227
228#[derive(Clone)]
259pub struct GuardLayer<G> {
260 guard: G,
261}
262
263impl<G> GuardLayer<G> {
264 pub fn new(guard: G) -> Self {
269 Self { guard }
270 }
271}
272
273impl<G, S> tower::Layer<S> for GuardLayer<G>
274where
275 G: Clone,
276{
277 type Service = GuardService<G, S>;
278
279 fn layer(&self, inner: S) -> Self::Service {
280 GuardService {
281 guard: self.guard.clone(),
282 inner,
283 }
284 }
285}
286
287#[doc(hidden)]
291#[derive(Clone)]
292pub struct GuardService<G, S> {
293 guard: G,
294 inner: S,
295}
296
297impl<G, S, R> Service<ToolRequest> for GuardService<G, S>
298where
299 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
300 S: Service<ToolRequest, Response = R> + Clone + Send + 'static,
301 S::Error: Into<Error> + Send,
302 S::Future: Send,
303 R: Send + 'static,
304{
305 type Response = R;
306 type Error = Error;
307 type Future = Pin<Box<dyn Future<Output = std::result::Result<R, Error>> + Send>>;
308
309 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
310 self.inner.poll_ready(cx).map_err(Into::into)
311 }
312
313 fn call(&mut self, req: ToolRequest) -> Self::Future {
314 match (self.guard)(&req) {
315 Ok(()) => {
316 let fut = self.inner.call(req);
317 Box::pin(async move { fut.await.map_err(Into::into) })
318 }
319 Err(msg) => Box::pin(async move { Err(Error::tool(msg)) }),
320 }
321 }
322}
323
324#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
344pub struct NoParams;
345
346impl<'de> serde::Deserialize<'de> for NoParams {
347 fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
348 where
349 D: serde::Deserializer<'de>,
350 {
351 struct NoParamsVisitor;
353
354 impl<'de> serde::de::Visitor<'de> for NoParamsVisitor {
355 type Value = NoParams;
356
357 fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
358 formatter.write_str("null or an object")
359 }
360
361 fn visit_unit<E>(self) -> std::result::Result<Self::Value, E>
362 where
363 E: serde::de::Error,
364 {
365 Ok(NoParams)
366 }
367
368 fn visit_none<E>(self) -> std::result::Result<Self::Value, E>
369 where
370 E: serde::de::Error,
371 {
372 Ok(NoParams)
373 }
374
375 fn visit_some<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
376 where
377 D: serde::Deserializer<'de>,
378 {
379 serde::Deserialize::deserialize(deserializer)
380 }
381
382 fn visit_map<A>(self, mut map: A) -> std::result::Result<Self::Value, A::Error>
383 where
384 A: serde::de::MapAccess<'de>,
385 {
386 while map
388 .next_entry::<serde::de::IgnoredAny, serde::de::IgnoredAny>()?
389 .is_some()
390 {}
391 Ok(NoParams)
392 }
393 }
394
395 deserializer.deserialize_any(NoParamsVisitor)
396 }
397}
398
399impl JsonSchema for NoParams {
400 fn schema_name() -> Cow<'static, str> {
401 Cow::Borrowed("NoParams")
402 }
403
404 fn json_schema(_generator: &mut SchemaGenerator) -> Schema {
405 serde_json::json!({
406 "type": "object"
407 })
408 .try_into()
409 .expect("valid schema")
410 }
411}
412
413pub(crate) fn validate_tool_name(name: &str) -> Result<()> {
422 if name.is_empty() {
423 return Err(Error::tool("Tool name cannot be empty"));
424 }
425 if name.len() > 64 {
426 return Err(Error::tool(format!(
427 "Tool name '{}' exceeds maximum length of 64 characters (got {})",
428 name,
429 name.len()
430 )));
431 }
432 if let Some(invalid_char) = name
433 .chars()
434 .find(|c| !c.is_ascii_alphanumeric() && *c != '_' && *c != '-' && *c != '.' && *c != '/')
435 {
436 return Err(Error::tool(format!(
437 "Tool name '{}' contains invalid character '{}'. Only alphanumeric, underscore, hyphen, dot, and forward slash are allowed.",
438 name, invalid_char
439 )));
440 }
441 Ok(())
442}
443
444pub(crate) fn ensure_object_schema(mut schema: Value) -> Value {
450 if let Some(obj) = schema.as_object_mut()
451 && !obj.contains_key("type")
452 {
453 obj.insert("type".to_string(), serde_json::json!("object"));
454 }
455 schema
456}
457
458pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
460
461pub trait ToolHandler: Send + Sync {
463 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>>;
465
466 fn call_with_context(
471 &self,
472 _ctx: RequestContext,
473 args: Value,
474 ) -> BoxFuture<'_, Result<CallToolResult>> {
475 self.call(args)
476 }
477
478 fn uses_context(&self) -> bool {
480 false
481 }
482
483 fn input_schema(&self) -> Value;
485}
486
487#[cfg(feature = "stateless")]
490pub trait MrtrToolHandler: Send + Sync {
491 fn call(
493 &self,
494 ctx: RequestContext,
495 args: Value,
496 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>>;
497
498 fn input_schema(&self) -> Value;
500}
501
502#[cfg(feature = "stateless")]
504struct MrtrToolHandlerService<H> {
505 handler: Arc<H>,
506}
507
508#[cfg(feature = "stateless")]
509impl<H> MrtrToolHandlerService<H> {
510 fn new(handler: H) -> Self {
511 Self {
512 handler: Arc::new(handler),
513 }
514 }
515}
516
517#[cfg(feature = "stateless")]
518impl<H> Clone for MrtrToolHandlerService<H> {
519 fn clone(&self) -> Self {
520 Self {
521 handler: self.handler.clone(),
522 }
523 }
524}
525
526#[cfg(feature = "stateless")]
527impl<H> Service<ToolRequest> for MrtrToolHandlerService<H>
528where
529 H: MrtrToolHandler + 'static,
530{
531 type Response = RequestOutcome<CallToolResult>;
532 type Error = Error;
533 type Future =
534 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
535
536 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
537 Poll::Ready(Ok(()))
538 }
539
540 fn call(&mut self, req: ToolRequest) -> Self::Future {
541 let handler = self.handler.clone();
542 Box::pin(async move { handler.call(req.ctx, req.args).await })
543 }
544}
545
546#[cfg(feature = "stateless")]
548struct ServiceMrtrToolHandler {
549 service: Mutex<BoxMrtrToolService>,
550 input_schema: Value,
551}
552
553#[cfg(feature = "stateless")]
554struct GuardedMrtrToolHandler<G> {
555 guard: G,
556 inner: Arc<dyn MrtrToolHandler>,
557}
558
559#[cfg(feature = "stateless")]
560impl<G> MrtrToolHandler for GuardedMrtrToolHandler<G>
561where
562 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
563{
564 fn call(
565 &self,
566 ctx: RequestContext,
567 args: Value,
568 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
569 let request = ToolRequest::new(ctx, args);
570 match (self.guard)(&request) {
571 Ok(()) => self.inner.call(request.ctx, request.args),
572 Err(message) => {
573 Box::pin(
574 async move { Ok(RequestOutcome::Complete(CallToolResult::error(message))) },
575 )
576 }
577 }
578 }
579
580 fn input_schema(&self) -> Value {
581 self.inner.input_schema()
582 }
583}
584
585#[cfg(feature = "stateless")]
586impl MrtrToolHandler for ServiceMrtrToolHandler {
587 fn call(
588 &self,
589 ctx: RequestContext,
590 args: Value,
591 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
592 Box::pin(async move {
593 let mut service = self.service.lock().await.clone();
594 let outcome = service
595 .ready()
596 .await
597 .expect("MRTR tool service is infallible")
598 .call(ToolRequest::new(ctx, args))
599 .await
600 .expect("MRTR tool service is infallible");
601 Ok(outcome)
602 })
603 }
604
605 fn input_schema(&self) -> Value {
606 self.input_schema.clone()
607 }
608}
609
610pub(crate) struct ToolHandlerService<H> {
615 handler: Arc<H>,
616}
617
618impl<H> ToolHandlerService<H> {
619 pub(crate) fn new(handler: H) -> Self {
620 Self {
621 handler: Arc::new(handler),
622 }
623 }
624}
625
626impl<H> Clone for ToolHandlerService<H> {
627 fn clone(&self) -> Self {
628 Self {
629 handler: self.handler.clone(),
630 }
631 }
632}
633
634impl<H> Service<ToolRequest> for ToolHandlerService<H>
635where
636 H: ToolHandler + 'static,
637{
638 type Response = CallToolResult;
639 type Error = Error;
640 type Future = Pin<Box<dyn Future<Output = std::result::Result<CallToolResult, Error>> + Send>>;
641
642 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
643 Poll::Ready(Ok(()))
644 }
645
646 fn call(&mut self, req: ToolRequest) -> Self::Future {
647 let handler = self.handler.clone();
648 Box::pin(async move { handler.call_with_context(req.ctx, req.args).await })
649 }
650}
651
652pub struct Tool {
659 pub name: String,
661 pub title: Option<String>,
663 pub description: Option<String>,
665 pub output_schema: Option<Value>,
667 pub icons: Option<Vec<ToolIcon>>,
669 pub annotations: Option<ToolAnnotations>,
671 pub meta: Option<Value>,
673 pub task_support: TaskSupportMode,
675 pub(crate) required_client_capabilities: Option<ClientCapabilities>,
678 pub(crate) service: Option<BoxToolService>,
680 #[cfg(feature = "stateless")]
681 pub(crate) mrtr_handler: Option<Arc<dyn MrtrToolHandler>>,
682 pub(crate) input_schema: Value,
684}
685
686impl std::fmt::Debug for Tool {
687 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
688 f.debug_struct("Tool")
689 .field("name", &self.name)
690 .field("title", &self.title)
691 .field("description", &self.description)
692 .field("output_schema", &self.output_schema)
693 .field("icons", &self.icons)
694 .field("annotations", &self.annotations)
695 .field("meta", &self.meta)
696 .field("task_support", &self.task_support)
697 .field(
698 "required_client_capabilities",
699 &self.required_client_capabilities,
700 )
701 .finish_non_exhaustive()
702 }
703}
704
705unsafe impl Send for Tool {}
708unsafe impl Sync for Tool {}
709
710impl Clone for Tool {
711 fn clone(&self) -> Self {
712 Self {
713 name: self.name.clone(),
714 title: self.title.clone(),
715 description: self.description.clone(),
716 output_schema: self.output_schema.clone(),
717 icons: self.icons.clone(),
718 annotations: self.annotations.clone(),
719 meta: self.meta.clone(),
720 task_support: self.task_support,
721 required_client_capabilities: self.required_client_capabilities.clone(),
722 service: self.service.clone(),
723 #[cfg(feature = "stateless")]
724 mrtr_handler: self.mrtr_handler.clone(),
725 input_schema: self.input_schema.clone(),
726 }
727 }
728}
729
730impl Tool {
731 pub fn builder(name: impl Into<String>) -> ToolBuilder {
733 ToolBuilder::new(name)
734 }
735
736 pub fn definition(&self) -> ToolDefinition {
738 let execution = match self.task_support {
739 TaskSupportMode::Forbidden => None,
740 mode => Some(ToolExecution {
741 task_support: Some(mode),
742 }),
743 };
744 ToolDefinition {
745 name: self.name.clone(),
746 title: self.title.clone(),
747 description: self.description.clone(),
748 input_schema: self.input_schema.clone(),
749 output_schema: self.output_schema.clone(),
750 icons: self.icons.clone(),
751 annotations: self.annotations.clone(),
752 execution,
753 meta: self.meta.clone(),
754 }
755 }
756
757 pub fn with_meta(
759 mut self,
760 meta: Value,
761 ) -> std::result::Result<Self, crate::protocol::MetaValidationError> {
762 crate::protocol::validate_meta_object(&meta)?;
763 self.meta = Some(meta);
764 Ok(self)
765 }
766
767 pub fn call(&self, args: Value) -> BoxFuture<'static, CallToolResult> {
772 let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
773 self.call_with_context(ctx, args)
774 }
775
776 pub fn call_with_context(
787 &self,
788 ctx: RequestContext,
789 args: Value,
790 ) -> BoxFuture<'static, CallToolResult> {
791 let tool = self.clone();
792 Box::pin(async move {
793 match tool.call_outcome_with_context(ctx, args).await {
794 Ok(RequestOutcome::Complete(result)) => result,
795 Ok(RequestOutcome::InputRequired(_)) => CallToolResult::error(
796 "tool requires additional client input; use call_outcome_with_context",
797 ),
798 Err(error) => CallToolResult::error(error.to_string()),
799 }
800 })
801 }
802
803 pub fn call_outcome(
805 &self,
806 args: Value,
807 ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
808 let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
809 self.call_outcome_with_context(ctx, args)
810 }
811
812 pub fn call_outcome_with_context(
815 &self,
816 ctx: RequestContext,
817 args: Value,
818 ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
819 use tower::ServiceExt;
820 #[cfg(feature = "stateless")]
821 if let Some(handler) = self.mrtr_handler.clone() {
822 return Box::pin(async move { handler.call(ctx, args).await });
823 }
824 let service = self
825 .service
826 .clone()
827 .expect("tool must have a complete or MRTR handler");
828 Box::pin(async move {
829 let result = service.oneshot(ToolRequest::new(ctx, args)).await.unwrap();
830 Ok(RequestOutcome::Complete(result))
831 })
832 }
833
834 pub fn require_client_capabilities(mut self, required: ClientCapabilities) -> Self {
842 self.required_client_capabilities = Some(required);
843 self
844 }
845
846 pub fn required_client_capabilities(&self) -> Option<&ClientCapabilities> {
848 self.required_client_capabilities.as_ref()
849 }
850
851 pub fn with_guard<G>(self, guard: G) -> Self
880 where
881 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
882 {
883 #[cfg(feature = "stateless")]
884 if let Some(inner) = self.mrtr_handler.clone() {
885 return Tool {
886 mrtr_handler: Some(Arc::new(GuardedMrtrToolHandler { guard, inner })),
887 ..self
888 };
889 }
890
891 let guarded = GuardService {
892 guard,
893 inner: self
894 .service
895 .expect("tool must have a complete or MRTR handler"),
896 };
897 let caught = ToolCatchError::new(guarded);
898 Tool {
899 service: Some(BoxCloneService::new(caught)),
900 ..self
901 }
902 }
903
904 pub fn with_name_prefix(&self, prefix: &str) -> Self {
931 Self {
932 name: format!("{}.{}", prefix, self.name),
933 title: self.title.clone(),
934 description: self.description.clone(),
935 output_schema: self.output_schema.clone(),
936 icons: self.icons.clone(),
937 annotations: self.annotations.clone(),
938 meta: self.meta.clone(),
939 task_support: self.task_support,
940 required_client_capabilities: self.required_client_capabilities.clone(),
941 service: self.service.clone(),
942 #[cfg(feature = "stateless")]
943 mrtr_handler: self.mrtr_handler.clone(),
944 input_schema: self.input_schema.clone(),
945 }
946 }
947
948 #[allow(clippy::too_many_arguments)]
950 fn from_handler<H: ToolHandler + 'static>(
951 name: String,
952 title: Option<String>,
953 description: Option<String>,
954 output_schema: Option<Value>,
955 icons: Option<Vec<ToolIcon>>,
956 annotations: Option<ToolAnnotations>,
957 task_support: TaskSupportMode,
958 input_schema_override: Option<Value>,
959 handler: H,
960 ) -> Self {
961 let input_schema =
962 ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
963 let handler_service = ToolHandlerService::new(handler);
964 let catch_error = ToolCatchError::new(handler_service);
965 let service = BoxCloneService::new(catch_error);
966
967 Self {
968 name,
969 title,
970 description,
971 output_schema,
972 icons,
973 annotations,
974 meta: None,
975 task_support,
976 required_client_capabilities: None,
977 service: Some(service),
978 #[cfg(feature = "stateless")]
979 mrtr_handler: None,
980 input_schema,
981 }
982 }
983
984 #[cfg(feature = "stateless")]
985 #[allow(clippy::too_many_arguments)]
986 fn from_mrtr_handler<H: MrtrToolHandler + 'static>(
987 name: String,
988 title: Option<String>,
989 description: Option<String>,
990 output_schema: Option<Value>,
991 icons: Option<Vec<ToolIcon>>,
992 annotations: Option<ToolAnnotations>,
993 task_support: TaskSupportMode,
994 input_schema_override: Option<Value>,
995 handler: H,
996 ) -> Self {
997 let input_schema =
998 ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
999 Self {
1000 name,
1001 title,
1002 description,
1003 output_schema,
1004 icons,
1005 annotations,
1006 meta: None,
1007 task_support,
1008 required_client_capabilities: None,
1009 service: None,
1010 mrtr_handler: Some(Arc::new(handler)),
1011 input_schema,
1012 }
1013 }
1014}
1015
1016pub struct ToolBuilder {
1044 name: String,
1045 title: Option<String>,
1046 description: Option<String>,
1047 output_schema: Option<Value>,
1048 input_schema_override: Option<Value>,
1049 icons: Option<Vec<ToolIcon>>,
1050 annotations: Option<ToolAnnotations>,
1051 task_support: TaskSupportMode,
1052}
1053
1054impl ToolBuilder {
1055 pub fn new(name: impl Into<String>) -> Self {
1068 let name = name.into();
1069 if let Err(e) = validate_tool_name(&name) {
1070 panic!("{e}");
1071 }
1072 Self {
1073 name,
1074 title: None,
1075 description: None,
1076 output_schema: None,
1077 input_schema_override: None,
1078 icons: None,
1079 annotations: None,
1080 task_support: TaskSupportMode::default(),
1081 }
1082 }
1083
1084 pub fn try_new(name: impl Into<String>) -> Result<Self> {
1090 let name = name.into();
1091 validate_tool_name(&name)?;
1092 Ok(Self {
1093 name,
1094 title: None,
1095 description: None,
1096 output_schema: None,
1097 input_schema_override: None,
1098 icons: None,
1099 annotations: None,
1100 task_support: TaskSupportMode::default(),
1101 })
1102 }
1103
1104 pub fn title(mut self, title: impl Into<String>) -> Self {
1120 self.title = Some(title.into());
1121 self
1122 }
1123
1124 pub fn output_schema(mut self, schema: Value) -> Self {
1126 self.output_schema = Some(schema);
1127 self
1128 }
1129
1130 pub fn input_schema(mut self, schema: Value) -> Self {
1183 self.input_schema_override = Some(schema);
1184 self
1185 }
1186
1187 pub fn icon(mut self, src: impl Into<String>) -> Self {
1189 self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1190 src: src.into(),
1191 mime_type: None,
1192 sizes: None,
1193 theme: None,
1194 });
1195 self
1196 }
1197
1198 pub fn icon_with_meta(
1200 mut self,
1201 src: impl Into<String>,
1202 mime_type: Option<String>,
1203 sizes: Option<Vec<String>>,
1204 ) -> Self {
1205 self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1206 src: src.into(),
1207 mime_type,
1208 sizes,
1209 theme: None,
1210 });
1211 self
1212 }
1213
1214 pub fn description(mut self, description: impl Into<String>) -> Self {
1216 self.description = Some(description.into());
1217 self
1218 }
1219
1220 pub fn read_only(mut self) -> Self {
1222 self.annotations
1223 .get_or_insert_with(ToolAnnotations::default)
1224 .read_only_hint = true;
1225 self
1226 }
1227
1228 pub fn non_destructive(mut self) -> Self {
1230 self.annotations
1231 .get_or_insert_with(ToolAnnotations::default)
1232 .destructive_hint = false;
1233 self
1234 }
1235
1236 pub fn destructive(mut self) -> Self {
1238 self.annotations
1239 .get_or_insert_with(ToolAnnotations::default)
1240 .destructive_hint = true;
1241 self
1242 }
1243
1244 pub fn idempotent(mut self) -> Self {
1246 self.annotations
1247 .get_or_insert_with(ToolAnnotations::default)
1248 .idempotent_hint = true;
1249 self
1250 }
1251
1252 pub fn read_only_safe(mut self) -> Self {
1258 let ann = self
1259 .annotations
1260 .get_or_insert_with(ToolAnnotations::default);
1261 ann.read_only_hint = true;
1262 ann.idempotent_hint = true;
1263 ann.destructive_hint = false;
1264 self
1265 }
1266
1267 pub fn annotations(mut self, annotations: ToolAnnotations) -> Self {
1269 self.annotations = Some(annotations);
1270 self
1271 }
1272
1273 pub fn task_support(mut self, mode: TaskSupportMode) -> Self {
1275 self.task_support = mode;
1276 self
1277 }
1278
1279 pub fn no_params_handler<F, Fut>(self, handler: F) -> ToolBuilderWithNoParamsHandler<F>
1297 where
1298 F: Fn() -> Fut + Send + Sync + 'static,
1299 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1300 {
1301 ToolBuilderWithNoParamsHandler {
1302 name: self.name,
1303 title: self.title,
1304 description: self.description,
1305 output_schema: self.output_schema,
1306 input_schema_override: self.input_schema_override,
1307 icons: self.icons,
1308 annotations: self.annotations,
1309 task_support: self.task_support,
1310 handler,
1311 }
1312 }
1313
1314 pub fn handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithHandler<I, F>
1357 where
1358 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1359 F: Fn(I) -> Fut + Send + Sync + 'static,
1360 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1361 {
1362 ToolBuilderWithHandler {
1363 name: self.name,
1364 title: self.title,
1365 description: self.description,
1366 output_schema: self.output_schema,
1367 input_schema_override: self.input_schema_override,
1368 icons: self.icons,
1369 annotations: self.annotations,
1370 task_support: self.task_support,
1371 handler,
1372 _phantom: std::marker::PhantomData,
1373 }
1374 }
1375
1376 #[cfg(feature = "stateless")]
1383 pub fn mrtr_handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithMrtrHandler<I, F>
1384 where
1385 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1386 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
1387 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
1388 {
1389 ToolBuilderWithMrtrHandler {
1390 name: self.name,
1391 title: self.title,
1392 description: self.description,
1393 output_schema: self.output_schema,
1394 input_schema_override: self.input_schema_override,
1395 icons: self.icons,
1396 annotations: self.annotations,
1397 task_support: self.task_support,
1398 handler,
1399 _phantom: std::marker::PhantomData,
1400 }
1401 }
1402
1403 pub fn extractor_handler<S, F, T>(
1497 self,
1498 state: S,
1499 handler: F,
1500 ) -> crate::extract::ToolBuilderWithExtractor<S, F, T>
1501 where
1502 S: Clone + Send + Sync + 'static,
1503 F: crate::extract::ExtractorHandler<S, T> + Clone,
1504 T: Send + Sync + 'static,
1505 {
1506 let input_schema = ensure_object_schema(
1507 self.input_schema_override
1508 .unwrap_or_else(|| F::input_schema()),
1509 );
1510 crate::extract::ToolBuilderWithExtractor {
1511 name: self.name,
1512 title: self.title,
1513 description: self.description,
1514 output_schema: self.output_schema,
1515 icons: self.icons,
1516 annotations: self.annotations,
1517 task_support: self.task_support,
1518 state,
1519 handler,
1520 input_schema,
1521 _phantom: std::marker::PhantomData,
1522 }
1523 }
1524
1525 #[deprecated(
1559 since = "0.8.0",
1560 note = "Use `extractor_handler` instead -- it auto-detects JSON schema from `Json<T>` extractors without requiring a turbofish"
1561 )]
1562 #[allow(deprecated)]
1563 pub fn extractor_handler_typed<S, F, T, I>(
1564 self,
1565 state: S,
1566 handler: F,
1567 ) -> crate::extract::ToolBuilderWithTypedExtractor<S, F, T, I>
1568 where
1569 S: Clone + Send + Sync + 'static,
1570 F: crate::extract::TypedExtractorHandler<S, T, I> + Clone,
1571 T: Send + Sync + 'static,
1572 I: schemars::JsonSchema + Send + Sync + 'static,
1573 {
1574 crate::extract::ToolBuilderWithTypedExtractor {
1575 name: self.name,
1576 title: self.title,
1577 description: self.description,
1578 output_schema: self.output_schema,
1579 input_schema_override: self.input_schema_override,
1580 icons: self.icons,
1581 annotations: self.annotations,
1582 task_support: self.task_support,
1583 state,
1584 handler,
1585 _phantom: std::marker::PhantomData,
1586 }
1587 }
1588}
1589
1590struct NoParamsTypedHandler<F> {
1594 handler: F,
1595}
1596
1597impl<F, Fut> ToolHandler for NoParamsTypedHandler<F>
1598where
1599 F: Fn() -> Fut + Send + Sync + 'static,
1600 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1601{
1602 fn call(&self, _args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
1603 Box::pin(async move { (self.handler)().await })
1604 }
1605
1606 fn input_schema(&self) -> Value {
1607 serde_json::json!({ "type": "object" })
1608 }
1609}
1610
1611#[doc(hidden)]
1613pub struct ToolBuilderWithHandler<I, F> {
1614 name: String,
1615 title: Option<String>,
1616 description: Option<String>,
1617 output_schema: Option<Value>,
1618 input_schema_override: Option<Value>,
1619 icons: Option<Vec<ToolIcon>>,
1620 annotations: Option<ToolAnnotations>,
1621 task_support: TaskSupportMode,
1622 handler: F,
1623 _phantom: std::marker::PhantomData<I>,
1624}
1625
1626#[cfg(feature = "stateless")]
1628#[doc(hidden)]
1629pub struct ToolBuilderWithMrtrHandler<I, F> {
1630 name: String,
1631 title: Option<String>,
1632 description: Option<String>,
1633 output_schema: Option<Value>,
1634 input_schema_override: Option<Value>,
1635 icons: Option<Vec<ToolIcon>>,
1636 annotations: Option<ToolAnnotations>,
1637 task_support: TaskSupportMode,
1638 handler: F,
1639 _phantom: std::marker::PhantomData<I>,
1640}
1641
1642#[cfg(feature = "stateless")]
1644#[doc(hidden)]
1645pub struct ToolBuilderWithMrtrLayer<I, F, L> {
1646 name: String,
1647 title: Option<String>,
1648 description: Option<String>,
1649 output_schema: Option<Value>,
1650 input_schema_override: Option<Value>,
1651 icons: Option<Vec<ToolIcon>>,
1652 annotations: Option<ToolAnnotations>,
1653 task_support: TaskSupportMode,
1654 handler: F,
1655 layer: L,
1656 _phantom: std::marker::PhantomData<I>,
1657}
1658
1659#[doc(hidden)]
1663pub struct ToolBuilderWithNoParamsHandler<F> {
1664 name: String,
1665 title: Option<String>,
1666 description: Option<String>,
1667 output_schema: Option<Value>,
1668 input_schema_override: Option<Value>,
1669 icons: Option<Vec<ToolIcon>>,
1670 annotations: Option<ToolAnnotations>,
1671 task_support: TaskSupportMode,
1672 handler: F,
1673}
1674
1675impl<F, Fut> ToolBuilderWithNoParamsHandler<F>
1676where
1677 F: Fn() -> Fut + Send + Sync + 'static,
1678 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1679{
1680 pub fn build(self) -> Tool {
1682 Tool::from_handler(
1683 self.name,
1684 self.title,
1685 self.description,
1686 self.output_schema,
1687 self.icons,
1688 self.annotations,
1689 self.task_support,
1690 self.input_schema_override,
1691 NoParamsTypedHandler {
1692 handler: self.handler,
1693 },
1694 )
1695 }
1696
1697 pub fn layer<L>(self, layer: L) -> ToolBuilderWithNoParamsHandlerLayer<F, L> {
1701 ToolBuilderWithNoParamsHandlerLayer {
1702 name: self.name,
1703 title: self.title,
1704 description: self.description,
1705 output_schema: self.output_schema,
1706 input_schema_override: self.input_schema_override,
1707 icons: self.icons,
1708 annotations: self.annotations,
1709 task_support: self.task_support,
1710 handler: self.handler,
1711 layer,
1712 }
1713 }
1714
1715 pub fn guard<G>(self, guard: G) -> ToolBuilderWithNoParamsHandlerLayer<F, GuardLayer<G>>
1719 where
1720 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1721 {
1722 self.layer(GuardLayer::new(guard))
1723 }
1724}
1725
1726#[doc(hidden)]
1728pub struct ToolBuilderWithNoParamsHandlerLayer<F, L> {
1729 name: String,
1730 title: Option<String>,
1731 description: Option<String>,
1732 output_schema: Option<Value>,
1733 input_schema_override: Option<Value>,
1734 icons: Option<Vec<ToolIcon>>,
1735 annotations: Option<ToolAnnotations>,
1736 task_support: TaskSupportMode,
1737 handler: F,
1738 layer: L,
1739}
1740
1741#[allow(private_bounds)]
1742impl<F, Fut, L> ToolBuilderWithNoParamsHandlerLayer<F, L>
1743where
1744 F: Fn() -> Fut + Send + Sync + 'static,
1745 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1746 L: tower::Layer<ToolHandlerService<NoParamsTypedHandler<F>>> + Clone + Send + Sync + 'static,
1747 L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
1748 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
1749 <L::Service as Service<ToolRequest>>::Future: Send,
1750{
1751 pub fn build(self) -> Tool {
1753 let input_schema = ensure_object_schema(
1754 self.input_schema_override
1755 .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
1756 );
1757
1758 let handler_service = ToolHandlerService::new(NoParamsTypedHandler {
1759 handler: self.handler,
1760 });
1761 let layered = self.layer.layer(handler_service);
1762 let catch_error = ToolCatchError::new(layered);
1763 let service = BoxCloneService::new(catch_error);
1764
1765 Tool {
1766 name: self.name,
1767 title: self.title,
1768 description: self.description,
1769 output_schema: self.output_schema,
1770 icons: self.icons,
1771 annotations: self.annotations,
1772 meta: None,
1773 task_support: self.task_support,
1774 required_client_capabilities: None,
1775 service: Some(service),
1776 #[cfg(feature = "stateless")]
1777 mrtr_handler: None,
1778 input_schema,
1779 }
1780 }
1781
1782 pub fn layer<L2>(
1784 self,
1785 layer: L2,
1786 ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<L2, L>> {
1787 ToolBuilderWithNoParamsHandlerLayer {
1788 name: self.name,
1789 title: self.title,
1790 description: self.description,
1791 output_schema: self.output_schema,
1792 input_schema_override: self.input_schema_override,
1793 icons: self.icons,
1794 annotations: self.annotations,
1795 task_support: self.task_support,
1796 handler: self.handler,
1797 layer: tower::layer::util::Stack::new(layer, self.layer),
1798 }
1799 }
1800
1801 pub fn guard<G>(
1805 self,
1806 guard: G,
1807 ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<GuardLayer<G>, L>>
1808 where
1809 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1810 {
1811 self.layer(GuardLayer::new(guard))
1812 }
1813}
1814
1815impl<I, F, Fut> ToolBuilderWithHandler<I, F>
1816where
1817 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1818 F: Fn(I) -> Fut + Send + Sync + 'static,
1819 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1820{
1821 pub fn build(self) -> Tool {
1823 Tool::from_handler(
1824 self.name,
1825 self.title,
1826 self.description,
1827 self.output_schema,
1828 self.icons,
1829 self.annotations,
1830 self.task_support,
1831 self.input_schema_override,
1832 TypedHandler {
1833 handler: self.handler,
1834 _phantom: std::marker::PhantomData,
1835 },
1836 )
1837 }
1838
1839 pub fn layer<L>(self, layer: L) -> ToolBuilderWithLayer<I, F, L> {
1865 ToolBuilderWithLayer {
1866 name: self.name,
1867 title: self.title,
1868 description: self.description,
1869 output_schema: self.output_schema,
1870 input_schema_override: self.input_schema_override,
1871 icons: self.icons,
1872 annotations: self.annotations,
1873 task_support: self.task_support,
1874 handler: self.handler,
1875 layer,
1876 _phantom: std::marker::PhantomData,
1877 }
1878 }
1879
1880 pub fn guard<G>(self, guard: G) -> ToolBuilderWithLayer<I, F, GuardLayer<G>>
1887 where
1888 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1889 {
1890 self.layer(GuardLayer::new(guard))
1891 }
1892}
1893
1894#[cfg(feature = "stateless")]
1895impl<I, F, Fut> ToolBuilderWithMrtrHandler<I, F>
1896where
1897 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1898 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
1899 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
1900{
1901 pub fn build(self) -> Tool {
1903 Tool::from_mrtr_handler(
1904 self.name,
1905 self.title,
1906 self.description,
1907 self.output_schema,
1908 self.icons,
1909 self.annotations,
1910 self.task_support,
1911 self.input_schema_override,
1912 TypedMrtrHandler {
1913 handler: self.handler,
1914 _phantom: std::marker::PhantomData,
1915 },
1916 )
1917 }
1918
1919 pub fn layer<L>(self, layer: L) -> ToolBuilderWithMrtrLayer<I, F, L> {
1925 ToolBuilderWithMrtrLayer {
1926 name: self.name,
1927 title: self.title,
1928 description: self.description,
1929 output_schema: self.output_schema,
1930 input_schema_override: self.input_schema_override,
1931 icons: self.icons,
1932 annotations: self.annotations,
1933 task_support: self.task_support,
1934 handler: self.handler,
1935 layer,
1936 _phantom: std::marker::PhantomData,
1937 }
1938 }
1939
1940 pub fn guard<G>(self, guard: G) -> ToolBuilderWithMrtrLayer<I, F, GuardLayer<G>>
1942 where
1943 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1944 {
1945 self.layer(GuardLayer::new(guard))
1946 }
1947}
1948
1949#[cfg(feature = "stateless")]
1950#[allow(private_bounds)]
1951impl<I, F, Fut, L> ToolBuilderWithMrtrLayer<I, F, L>
1952where
1953 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1954 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
1955 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
1956 L: tower::Layer<MrtrToolHandlerService<TypedMrtrHandler<I, F>>> + Clone + Send + Sync + 'static,
1957 L::Service:
1958 Service<ToolRequest, Response = RequestOutcome<CallToolResult>> + Clone + Send + 'static,
1959 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send + 'static,
1960 <L::Service as Service<ToolRequest>>::Future: Send + 'static,
1961{
1962 pub fn build(self) -> Tool {
1964 let input_schema = self.input_schema_override.unwrap_or_else(|| {
1965 let schema = schemars::schema_for!(I);
1966 serde_json::to_value(schema).unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
1967 });
1968 let input_schema = ensure_object_schema(input_schema);
1969 let service = MrtrToolHandlerService::new(TypedMrtrHandler {
1970 handler: self.handler,
1971 _phantom: std::marker::PhantomData,
1972 });
1973 let service = self.layer.layer(service);
1974 let service = BoxCloneService::new(MrtrToolCatchError::new(service));
1975
1976 Tool {
1977 name: self.name,
1978 title: self.title,
1979 description: self.description,
1980 output_schema: self.output_schema,
1981 icons: self.icons,
1982 annotations: self.annotations,
1983 meta: None,
1984 task_support: self.task_support,
1985 required_client_capabilities: None,
1986 service: None,
1987 mrtr_handler: Some(Arc::new(ServiceMrtrToolHandler {
1988 service: Mutex::new(service),
1989 input_schema: input_schema.clone(),
1990 })),
1991 input_schema,
1992 }
1993 }
1994
1995 pub fn layer<L2>(
1997 self,
1998 layer: L2,
1999 ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<L2, L>> {
2000 ToolBuilderWithMrtrLayer {
2001 name: self.name,
2002 title: self.title,
2003 description: self.description,
2004 output_schema: self.output_schema,
2005 input_schema_override: self.input_schema_override,
2006 icons: self.icons,
2007 annotations: self.annotations,
2008 task_support: self.task_support,
2009 handler: self.handler,
2010 layer: tower::layer::util::Stack::new(layer, self.layer),
2011 _phantom: std::marker::PhantomData,
2012 }
2013 }
2014
2015 pub fn guard<G>(
2017 self,
2018 guard: G,
2019 ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2020 where
2021 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2022 {
2023 self.layer(GuardLayer::new(guard))
2024 }
2025}
2026
2027#[doc(hidden)]
2031pub struct ToolBuilderWithLayer<I, F, L> {
2032 name: String,
2033 title: Option<String>,
2034 description: Option<String>,
2035 output_schema: Option<Value>,
2036 input_schema_override: Option<Value>,
2037 icons: Option<Vec<ToolIcon>>,
2038 annotations: Option<ToolAnnotations>,
2039 task_support: TaskSupportMode,
2040 handler: F,
2041 layer: L,
2042 _phantom: std::marker::PhantomData<I>,
2043}
2044
2045#[allow(private_bounds)]
2048impl<I, F, Fut, L> ToolBuilderWithLayer<I, F, L>
2049where
2050 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2051 F: Fn(I) -> Fut + Send + Sync + 'static,
2052 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2053 L: tower::Layer<ToolHandlerService<TypedHandler<I, F>>> + Clone + Send + Sync + 'static,
2054 L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
2055 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
2056 <L::Service as Service<ToolRequest>>::Future: Send,
2057{
2058 pub fn build(self) -> Tool {
2060 let input_schema = self.input_schema_override.unwrap_or_else(|| {
2061 let input_schema = schemars::schema_for!(I);
2062 serde_json::to_value(input_schema)
2063 .unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
2064 });
2065 let input_schema = ensure_object_schema(input_schema);
2066
2067 let handler_service = ToolHandlerService::new(TypedHandler {
2068 handler: self.handler,
2069 _phantom: std::marker::PhantomData,
2070 });
2071 let layered = self.layer.layer(handler_service);
2072 let catch_error = ToolCatchError::new(layered);
2073 let service = BoxCloneService::new(catch_error);
2074
2075 Tool {
2076 name: self.name,
2077 title: self.title,
2078 description: self.description,
2079 output_schema: self.output_schema,
2080 icons: self.icons,
2081 annotations: self.annotations,
2082 meta: None,
2083 task_support: self.task_support,
2084 required_client_capabilities: None,
2085 service: Some(service),
2086 #[cfg(feature = "stateless")]
2087 mrtr_handler: None,
2088 input_schema,
2089 }
2090 }
2091
2092 pub fn layer<L2>(
2097 self,
2098 layer: L2,
2099 ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<L2, L>> {
2100 ToolBuilderWithLayer {
2101 name: self.name,
2102 title: self.title,
2103 description: self.description,
2104 output_schema: self.output_schema,
2105 input_schema_override: self.input_schema_override,
2106 icons: self.icons,
2107 annotations: self.annotations,
2108 task_support: self.task_support,
2109 handler: self.handler,
2110 layer: tower::layer::util::Stack::new(layer, self.layer),
2111 _phantom: std::marker::PhantomData,
2112 }
2113 }
2114
2115 pub fn guard<G>(
2119 self,
2120 guard: G,
2121 ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2122 where
2123 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2124 {
2125 self.layer(GuardLayer::new(guard))
2126 }
2127}
2128
2129struct TypedHandler<I, F> {
2135 handler: F,
2136 _phantom: std::marker::PhantomData<I>,
2137}
2138
2139#[cfg(feature = "stateless")]
2140struct TypedMrtrHandler<I, F> {
2141 handler: F,
2142 _phantom: std::marker::PhantomData<I>,
2143}
2144
2145#[cfg(feature = "stateless")]
2146impl<I, F, Fut> MrtrToolHandler for TypedMrtrHandler<I, F>
2147where
2148 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2149 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2150 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2151{
2152 fn call(
2153 &self,
2154 ctx: RequestContext,
2155 args: Value,
2156 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
2157 Box::pin(async move {
2158 let input: I = serde_json::from_value(args)
2159 .map_err(|error| Error::invalid_params(format!("Invalid input: {error}")))?;
2160 (self.handler)(ctx, input).await
2161 })
2162 }
2163
2164 fn input_schema(&self) -> Value {
2165 let schema = schemars::schema_for!(I);
2166 ensure_object_schema(
2167 serde_json::to_value(schema)
2168 .unwrap_or_else(|_| serde_json::json!({ "type": "object" })),
2169 )
2170 }
2171}
2172
2173impl<I, F, Fut> ToolHandler for TypedHandler<I, F>
2174where
2175 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2176 F: Fn(I) -> Fut + Send + Sync + 'static,
2177 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2178{
2179 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2180 Box::pin(async move {
2181 let input: I = match serde_json::from_value(args) {
2182 Ok(input) => input,
2183 Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2184 };
2185 (self.handler)(input).await
2186 })
2187 }
2188
2189 fn input_schema(&self) -> Value {
2190 let schema = schemars::schema_for!(I);
2191 let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2192 serde_json::json!({
2193 "type": "object"
2194 })
2195 });
2196 ensure_object_schema(schema)
2197 }
2198}
2199
2200pub trait McpTool: Send + Sync + 'static {
2241 const NAME: &'static str;
2243 const DESCRIPTION: &'static str;
2245
2246 type Input: JsonSchema + DeserializeOwned + Send;
2248 type Output: Serialize + Send;
2250
2251 fn call(&self, input: Self::Input) -> impl Future<Output = Result<Self::Output>> + Send;
2253
2254 fn annotations(&self) -> Option<ToolAnnotations> {
2256 None
2257 }
2258
2259 fn into_tool(self) -> Tool
2267 where
2268 Self: Sized,
2269 {
2270 if let Err(e) = validate_tool_name(Self::NAME) {
2271 panic!("{e}");
2272 }
2273 let annotations = self.annotations();
2274 let tool = Arc::new(self);
2275 Tool::from_handler(
2276 Self::NAME.to_string(),
2277 None,
2278 Some(Self::DESCRIPTION.to_string()),
2279 None,
2280 None,
2281 annotations,
2282 TaskSupportMode::default(),
2283 None,
2284 McpToolHandler { tool },
2285 )
2286 }
2287}
2288
2289struct McpToolHandler<T: McpTool> {
2291 tool: Arc<T>,
2292}
2293
2294impl<T: McpTool> ToolHandler for McpToolHandler<T> {
2295 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2296 let tool = self.tool.clone();
2297 Box::pin(async move {
2298 let input: T::Input = match serde_json::from_value(args) {
2299 Ok(input) => input,
2300 Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2301 };
2302 let output = tool.call(input).await?;
2303 let value = serde_json::to_value(output).tool_context("Failed to serialize output")?;
2304 Ok(CallToolResult::json(value))
2305 })
2306 }
2307
2308 fn input_schema(&self) -> Value {
2309 let schema = schemars::schema_for!(T::Input);
2310 let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2311 serde_json::json!({
2312 "type": "object"
2313 })
2314 });
2315 ensure_object_schema(schema)
2316 }
2317}
2318
2319#[cfg(test)]
2320mod tests {
2321 use super::*;
2322 use crate::extract::{Context, Json, RawArgs, State};
2323 use crate::protocol::Content;
2324 use schemars::JsonSchema;
2325 use serde::Deserialize;
2326
2327 #[derive(Debug, Deserialize, JsonSchema)]
2328 struct GreetInput {
2329 name: String,
2330 }
2331
2332 #[tokio::test]
2333 async fn test_builder_tool() {
2334 let tool = ToolBuilder::new("greet")
2335 .description("Greet someone")
2336 .handler(|input: GreetInput| async move {
2337 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2338 })
2339 .build();
2340
2341 assert_eq!(tool.name, "greet");
2342 assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2343
2344 let result = tool.call(serde_json::json!({"name": "World"})).await;
2345
2346 assert!(!result.is_error);
2347 }
2348
2349 #[cfg(feature = "stateless")]
2350 #[tokio::test]
2351 async fn test_mrtr_builder_preserves_input_required_outcome() {
2352 let tool = ToolBuilder::new("continue")
2353 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2354 Ok(RequestOutcome::input_required(
2355 crate::protocol::InputRequiredResult::new().with_request_state("signed-state"),
2356 ))
2357 })
2358 .build();
2359
2360 let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2361 assert_eq!(
2362 outcome
2363 .as_input_required()
2364 .and_then(|result| result.request_state.as_deref()),
2365 Some("signed-state")
2366 );
2367 }
2368
2369 #[cfg(feature = "stateless")]
2370 #[tokio::test]
2371 async fn mrtr_builder_composes_guards_and_layers() {
2372 use std::sync::atomic::{AtomicUsize, Ordering};
2373 use std::time::Duration;
2374 use tower::timeout::TimeoutLayer;
2375
2376 let rounds = Arc::new(AtomicUsize::new(0));
2377 let observed = rounds.clone();
2378 let tool = ToolBuilder::new("guarded_continue")
2379 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2380 Ok(RequestOutcome::input_required(
2381 crate::protocol::InputRequiredResult::new().with_request_state("continue"),
2382 ))
2383 })
2384 .layer(TimeoutLayer::new(Duration::from_secs(1)))
2385 .guard(move |_request| {
2386 observed.fetch_add(1, Ordering::SeqCst);
2387 Ok(())
2388 })
2389 .build();
2390
2391 for _ in 0..2 {
2392 assert!(
2393 tool.call_outcome(serde_json::json!({}))
2394 .await
2395 .unwrap()
2396 .as_input_required()
2397 .is_some()
2398 );
2399 }
2400 assert_eq!(rounds.load(Ordering::SeqCst), 2);
2401 }
2402
2403 #[cfg(feature = "stateless")]
2404 #[tokio::test]
2405 async fn built_mrtr_tool_accepts_a_guard() {
2406 let tool = ToolBuilder::new("denied_continue")
2407 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2408 Ok(RequestOutcome::input_required(
2409 crate::protocol::InputRequiredResult::new().with_request_state("unreachable"),
2410 ))
2411 })
2412 .build()
2413 .with_guard(|_request| Err("MRTR access denied".to_string()));
2414
2415 let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2416 let result = outcome
2417 .as_complete()
2418 .expect("guard rejection is a complete tool error");
2419 assert!(result.is_error);
2420 assert_eq!(result.first_text(), Some("MRTR access denied"));
2421 }
2422
2423 #[tokio::test]
2424 async fn test_raw_handler() {
2425 let tool = ToolBuilder::new("echo")
2426 .description("Echo input")
2427 .extractor_handler((), |RawArgs(args): RawArgs| async move {
2428 Ok(CallToolResult::json(args))
2429 })
2430 .build();
2431
2432 let result = tool.call(serde_json::json!({"foo": "bar"})).await;
2433
2434 assert!(!result.is_error);
2435 }
2436
2437 #[test]
2438 fn test_invalid_tool_name_empty() {
2439 let err = ToolBuilder::try_new("").err().expect("should fail");
2440 assert!(err.to_string().contains("cannot be empty"));
2441 }
2442
2443 #[test]
2444 fn test_invalid_tool_name_too_long() {
2445 let long_name = "a".repeat(65);
2446 let err = ToolBuilder::try_new(long_name).err().expect("should fail");
2447 assert!(err.to_string().contains("exceeds maximum"));
2448 }
2449
2450 #[test]
2451 fn test_invalid_tool_name_bad_chars() {
2452 let err = ToolBuilder::try_new("my tool!").err().expect("should fail");
2453 assert!(err.to_string().contains("invalid character"));
2454 }
2455
2456 #[test]
2457 #[should_panic(expected = "cannot be empty")]
2458 fn test_new_panics_on_empty_name() {
2459 ToolBuilder::new("");
2460 }
2461
2462 #[test]
2463 #[should_panic(expected = "exceeds maximum")]
2464 fn test_new_panics_on_too_long_name() {
2465 ToolBuilder::new("a".repeat(65));
2466 }
2467
2468 #[test]
2469 #[should_panic(expected = "invalid character")]
2470 fn test_new_panics_on_invalid_chars() {
2471 ToolBuilder::new("my tool!");
2472 }
2473
2474 #[test]
2475 fn test_valid_tool_names() {
2476 let names = [
2478 "my_tool",
2479 "my-tool",
2480 "my.tool",
2481 "my/tool",
2482 "user-profile/update",
2483 "MyTool123",
2484 "a",
2485 &"a".repeat(64),
2486 ];
2487 for name in names {
2488 assert!(
2489 ToolBuilder::try_new(name).is_ok(),
2490 "Expected '{}' to be valid",
2491 name
2492 );
2493 }
2494 }
2495
2496 #[tokio::test]
2497 async fn test_context_aware_handler() {
2498 use crate::context::notification_channel;
2499 use crate::protocol::{ProgressToken, RequestId};
2500
2501 #[derive(Debug, Deserialize, JsonSchema)]
2502 struct ProcessInput {
2503 count: i32,
2504 }
2505
2506 let tool = ToolBuilder::new("process")
2507 .description("Process with context")
2508 .extractor_handler(
2509 (),
2510 |ctx: Context, Json(input): Json<ProcessInput>| async move {
2511 for i in 0..input.count {
2513 if ctx.is_cancelled() {
2514 return Ok(CallToolResult::error("Cancelled"));
2515 }
2516 ctx.report_progress(i as f64, Some(input.count as f64), None)
2517 .await;
2518 }
2519 Ok(CallToolResult::text(format!(
2520 "Processed {} items",
2521 input.count
2522 )))
2523 },
2524 )
2525 .build();
2526
2527 assert_eq!(tool.name, "process");
2528
2529 let (tx, mut rx) = notification_channel(10);
2531 let ctx = RequestContext::new(RequestId::Number(1))
2532 .with_progress_token(ProgressToken::Number(42))
2533 .with_notification_sender(tx);
2534
2535 let result = tool
2536 .call_with_context(ctx, serde_json::json!({"count": 3}))
2537 .await;
2538
2539 assert!(!result.is_error);
2540
2541 let mut progress_count = 0;
2543 while rx.try_recv().is_ok() {
2544 progress_count += 1;
2545 }
2546 assert_eq!(progress_count, 3);
2547 }
2548
2549 #[tokio::test]
2550 async fn test_context_aware_handler_cancellation() {
2551 use crate::protocol::RequestId;
2552 use std::sync::atomic::{AtomicI32, Ordering};
2553
2554 #[derive(Debug, Deserialize, JsonSchema)]
2555 struct LongRunningInput {
2556 iterations: i32,
2557 }
2558
2559 let iterations_completed = Arc::new(AtomicI32::new(0));
2560 let iterations_ref = iterations_completed.clone();
2561
2562 let tool = ToolBuilder::new("long_running")
2563 .description("Long running task")
2564 .extractor_handler(
2565 (),
2566 move |ctx: Context, Json(input): Json<LongRunningInput>| {
2567 let completed = iterations_ref.clone();
2568 async move {
2569 for i in 0..input.iterations {
2570 if ctx.is_cancelled() {
2571 return Ok(CallToolResult::error("Cancelled"));
2572 }
2573 completed.fetch_add(1, Ordering::SeqCst);
2574 tokio::task::yield_now().await;
2576 if i == 2 {
2578 ctx.cancellation_token().cancel();
2579 }
2580 }
2581 Ok(CallToolResult::text("Done"))
2582 }
2583 },
2584 )
2585 .build();
2586
2587 let ctx = RequestContext::new(RequestId::Number(1));
2588
2589 let result = tool
2590 .call_with_context(ctx, serde_json::json!({"iterations": 10}))
2591 .await;
2592
2593 assert!(result.is_error);
2596 assert_eq!(iterations_completed.load(Ordering::SeqCst), 3);
2597 }
2598
2599 #[tokio::test]
2600 async fn test_tool_builder_with_enhanced_fields() {
2601 let output_schema = serde_json::json!({
2602 "type": "object",
2603 "properties": {
2604 "greeting": {"type": "string"}
2605 }
2606 });
2607
2608 let tool = ToolBuilder::new("greet")
2609 .title("Greeting Tool")
2610 .description("Greet someone")
2611 .output_schema(output_schema.clone())
2612 .icon("https://example.com/icon.png")
2613 .icon_with_meta(
2614 "https://example.com/icon-large.png",
2615 Some("image/png".to_string()),
2616 Some(vec!["96x96".to_string()]),
2617 )
2618 .handler(|input: GreetInput| async move {
2619 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2620 })
2621 .build();
2622
2623 assert_eq!(tool.name, "greet");
2624 assert_eq!(tool.title.as_deref(), Some("Greeting Tool"));
2625 assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2626 assert_eq!(tool.output_schema, Some(output_schema));
2627 assert!(tool.icons.is_some());
2628 assert_eq!(tool.icons.as_ref().unwrap().len(), 2);
2629
2630 let def = tool.definition();
2632 assert_eq!(def.title.as_deref(), Some("Greeting Tool"));
2633 assert!(def.output_schema.is_some());
2634 assert!(def.icons.is_some());
2635 }
2636
2637 #[tokio::test]
2638 async fn test_handler_with_state() {
2639 let shared = Arc::new("shared-state".to_string());
2640
2641 let tool = ToolBuilder::new("stateful")
2642 .description("Uses shared state")
2643 .extractor_handler(
2644 shared,
2645 |State(state): State<Arc<String>>, Json(input): Json<GreetInput>| async move {
2646 Ok(CallToolResult::text(format!(
2647 "{}: Hello, {}!",
2648 state, input.name
2649 )))
2650 },
2651 )
2652 .build();
2653
2654 let result = tool.call(serde_json::json!({"name": "World"})).await;
2655 assert!(!result.is_error);
2656 }
2657
2658 #[tokio::test]
2659 async fn test_handler_with_state_and_context() {
2660 use crate::protocol::RequestId;
2661
2662 let shared = Arc::new(42_i32);
2663
2664 let tool =
2665 ToolBuilder::new("stateful_ctx")
2666 .description("Uses state and context")
2667 .extractor_handler(
2668 shared,
2669 |State(state): State<Arc<i32>>,
2670 _ctx: Context,
2671 Json(input): Json<GreetInput>| async move {
2672 Ok(CallToolResult::text(format!(
2673 "{}: Hello, {}!",
2674 state, input.name
2675 )))
2676 },
2677 )
2678 .build();
2679
2680 let ctx = RequestContext::new(RequestId::Number(1));
2681 let result = tool
2682 .call_with_context(ctx, serde_json::json!({"name": "World"}))
2683 .await;
2684 assert!(!result.is_error);
2685 }
2686
2687 #[tokio::test]
2688 async fn test_handler_no_params() {
2689 let tool = ToolBuilder::new("no_params")
2690 .description("Takes no parameters")
2691 .extractor_handler((), |Json(_): Json<NoParams>| async {
2692 Ok(CallToolResult::text("no params result"))
2693 })
2694 .build();
2695
2696 assert_eq!(tool.name, "no_params");
2697
2698 let result = tool.call(serde_json::json!({})).await;
2700 assert!(!result.is_error);
2701
2702 let result = tool.call(serde_json::json!({"unexpected": "value"})).await;
2704 assert!(!result.is_error);
2705
2706 let schema = tool.definition().input_schema;
2708 assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2709 }
2710
2711 #[tokio::test]
2712 async fn test_handler_with_state_no_params() {
2713 let shared = Arc::new("shared_value".to_string());
2714
2715 let tool = ToolBuilder::new("with_state_no_params")
2716 .description("Takes no parameters but has state")
2717 .extractor_handler(
2718 shared,
2719 |State(state): State<Arc<String>>, Json(_): Json<NoParams>| async move {
2720 Ok(CallToolResult::text(format!("state: {}", state)))
2721 },
2722 )
2723 .build();
2724
2725 assert_eq!(tool.name, "with_state_no_params");
2726
2727 let result = tool.call(serde_json::json!({})).await;
2729 assert!(!result.is_error);
2730 assert_eq!(result.first_text().unwrap(), "state: shared_value");
2731
2732 let schema = tool.definition().input_schema;
2734 assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2735 }
2736
2737 #[tokio::test]
2738 async fn test_handler_no_params_with_context() {
2739 let tool = ToolBuilder::new("no_params_with_context")
2740 .description("Takes no parameters but has context")
2741 .extractor_handler((), |_ctx: Context, Json(_): Json<NoParams>| async move {
2742 Ok(CallToolResult::text("context available"))
2743 })
2744 .build();
2745
2746 assert_eq!(tool.name, "no_params_with_context");
2747
2748 let result = tool.call(serde_json::json!({})).await;
2749 assert!(!result.is_error);
2750 assert_eq!(result.first_text().unwrap(), "context available");
2751 }
2752
2753 #[tokio::test]
2754 async fn test_handler_with_state_and_context_no_params() {
2755 let shared = Arc::new("shared".to_string());
2756
2757 let tool = ToolBuilder::new("state_context_no_params")
2758 .description("Has state and context, no params")
2759 .extractor_handler(
2760 shared,
2761 |State(state): State<Arc<String>>,
2762 _ctx: Context,
2763 Json(_): Json<NoParams>| async move {
2764 Ok(CallToolResult::text(format!("state: {}", state)))
2765 },
2766 )
2767 .build();
2768
2769 assert_eq!(tool.name, "state_context_no_params");
2770
2771 let result = tool.call(serde_json::json!({})).await;
2772 assert!(!result.is_error);
2773 assert_eq!(result.first_text().unwrap(), "state: shared");
2774 }
2775
2776 #[tokio::test]
2777 async fn test_raw_handler_with_state() {
2778 let prefix = Arc::new("prefix:".to_string());
2779
2780 let tool = ToolBuilder::new("raw_with_state")
2781 .description("Raw handler with state")
2782 .extractor_handler(
2783 prefix,
2784 |State(state): State<Arc<String>>, RawArgs(args): RawArgs| async move {
2785 Ok(CallToolResult::text(format!("{} {}", state, args)))
2786 },
2787 )
2788 .build();
2789
2790 assert_eq!(tool.name, "raw_with_state");
2791
2792 let result = tool.call(serde_json::json!({"key": "value"})).await;
2793 assert!(!result.is_error);
2794 assert!(result.first_text().unwrap().starts_with("prefix:"));
2795 }
2796
2797 #[tokio::test]
2798 async fn test_raw_handler_with_state_and_context() {
2799 let prefix = Arc::new("prefix:".to_string());
2800
2801 let tool = ToolBuilder::new("raw_state_context")
2802 .description("Raw handler with state and context")
2803 .extractor_handler(
2804 prefix,
2805 |State(state): State<Arc<String>>,
2806 _ctx: Context,
2807 RawArgs(args): RawArgs| async move {
2808 Ok(CallToolResult::text(format!("{} {}", state, args)))
2809 },
2810 )
2811 .build();
2812
2813 assert_eq!(tool.name, "raw_state_context");
2814
2815 let result = tool.call(serde_json::json!({"key": "value"})).await;
2816 assert!(!result.is_error);
2817 assert!(result.first_text().unwrap().starts_with("prefix:"));
2818 }
2819
2820 #[tokio::test]
2821 async fn test_tool_with_timeout_layer() {
2822 use std::time::Duration;
2823 use tower::timeout::TimeoutLayer;
2824
2825 #[derive(Debug, Deserialize, JsonSchema)]
2826 struct SlowInput {
2827 delay_ms: u64,
2828 }
2829
2830 let tool = ToolBuilder::new("slow_tool")
2832 .description("A slow tool")
2833 .handler(|input: SlowInput| async move {
2834 tokio::time::sleep(Duration::from_millis(input.delay_ms)).await;
2835 Ok(CallToolResult::text("completed"))
2836 })
2837 .layer(TimeoutLayer::new(Duration::from_millis(50)))
2838 .build();
2839
2840 let result = tool.call(serde_json::json!({"delay_ms": 10})).await;
2842 assert!(!result.is_error);
2843 assert_eq!(result.first_text().unwrap(), "completed");
2844
2845 let result = tool.call(serde_json::json!({"delay_ms": 200})).await;
2847 assert!(result.is_error);
2848 let msg = result.first_text().unwrap().to_lowercase();
2850 assert!(
2851 msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
2852 "Expected timeout error, got: {}",
2853 msg
2854 );
2855 }
2856
2857 #[tokio::test]
2858 async fn test_tool_with_concurrency_limit_layer() {
2859 use std::sync::atomic::{AtomicU32, Ordering};
2860 use std::time::Duration;
2861 use tower::limit::ConcurrencyLimitLayer;
2862
2863 #[derive(Debug, Deserialize, JsonSchema)]
2864 struct WorkInput {
2865 id: u32,
2866 }
2867
2868 let max_concurrent = Arc::new(AtomicU32::new(0));
2869 let current_concurrent = Arc::new(AtomicU32::new(0));
2870 let max_ref = max_concurrent.clone();
2871 let current_ref = current_concurrent.clone();
2872
2873 let tool = ToolBuilder::new("concurrent_tool")
2875 .description("A concurrent tool")
2876 .handler(move |input: WorkInput| {
2877 let max = max_ref.clone();
2878 let current = current_ref.clone();
2879 async move {
2880 let prev = current.fetch_add(1, Ordering::SeqCst);
2882 max.fetch_max(prev + 1, Ordering::SeqCst);
2883
2884 tokio::time::sleep(Duration::from_millis(50)).await;
2886
2887 current.fetch_sub(1, Ordering::SeqCst);
2888 Ok(CallToolResult::text(format!("completed {}", input.id)))
2889 }
2890 })
2891 .layer(ConcurrencyLimitLayer::new(2))
2892 .build();
2893
2894 let handles: Vec<_> = (0..4)
2896 .map(|i| {
2897 let t = tool.call(serde_json::json!({"id": i}));
2898 tokio::spawn(t)
2899 })
2900 .collect();
2901
2902 for handle in handles {
2903 let result = handle.await.unwrap();
2904 assert!(!result.is_error);
2905 }
2906
2907 assert!(max_concurrent.load(Ordering::SeqCst) <= 2);
2909 }
2910
2911 #[tokio::test]
2912 async fn test_tool_with_multiple_layers() {
2913 use std::time::Duration;
2914 use tower::limit::ConcurrencyLimitLayer;
2915 use tower::timeout::TimeoutLayer;
2916
2917 #[derive(Debug, Deserialize, JsonSchema)]
2918 struct Input {
2919 value: String,
2920 }
2921
2922 let tool = ToolBuilder::new("multi_layer_tool")
2924 .description("Tool with multiple layers")
2925 .handler(|input: Input| async move {
2926 Ok(CallToolResult::text(format!("processed: {}", input.value)))
2927 })
2928 .layer(TimeoutLayer::new(Duration::from_secs(5)))
2929 .layer(ConcurrencyLimitLayer::new(10))
2930 .build();
2931
2932 let result = tool.call(serde_json::json!({"value": "test"})).await;
2933 assert!(!result.is_error);
2934 assert_eq!(result.first_text().unwrap(), "processed: test");
2935 }
2936
2937 #[test]
2938 fn test_tool_catch_error_clone() {
2939 let tool = ToolBuilder::new("test")
2942 .description("test")
2943 .extractor_handler((), |RawArgs(_args): RawArgs| async {
2944 Ok(CallToolResult::text("ok"))
2945 })
2946 .build();
2947 let _clone = tool.call(serde_json::json!({}));
2949 }
2950
2951 #[test]
2952 fn test_tool_catch_error_debug() {
2953 #[derive(Debug, Clone)]
2957 struct DebugService;
2958
2959 impl Service<ToolRequest> for DebugService {
2960 type Response = CallToolResult;
2961 type Error = crate::error::Error;
2962 type Future = Pin<
2963 Box<
2964 dyn Future<Output = std::result::Result<CallToolResult, crate::error::Error>>
2965 + Send,
2966 >,
2967 >;
2968
2969 fn poll_ready(
2970 &mut self,
2971 _cx: &mut std::task::Context<'_>,
2972 ) -> Poll<std::result::Result<(), Self::Error>> {
2973 Poll::Ready(Ok(()))
2974 }
2975
2976 fn call(&mut self, _req: ToolRequest) -> Self::Future {
2977 Box::pin(async { Ok(CallToolResult::text("ok")) })
2978 }
2979 }
2980
2981 let catch_error = ToolCatchError::new(DebugService);
2982 let debug = format!("{:?}", catch_error);
2983 assert!(debug.contains("ToolCatchError"));
2984 }
2985
2986 #[test]
2987 fn test_tool_request_new() {
2988 use crate::protocol::RequestId;
2989
2990 let ctx = RequestContext::new(RequestId::Number(42));
2991 let args = serde_json::json!({"key": "value"});
2992 let req = ToolRequest::new(ctx.clone(), args.clone());
2993
2994 assert_eq!(req.args, args);
2995 }
2996
2997 #[test]
2998 fn test_no_params_schema() {
2999 let schema = schemars::schema_for!(NoParams);
3001 let schema_value = serde_json::to_value(&schema).unwrap();
3002 assert_eq!(
3003 schema_value.get("type").and_then(|v| v.as_str()),
3004 Some("object"),
3005 "NoParams should generate type: object schema"
3006 );
3007 }
3008
3009 #[test]
3010 fn test_no_params_deserialize() {
3011 let from_empty_object: NoParams = serde_json::from_str("{}").unwrap();
3013 assert_eq!(from_empty_object, NoParams);
3014
3015 let from_null: NoParams = serde_json::from_str("null").unwrap();
3016 assert_eq!(from_null, NoParams);
3017
3018 let from_object_with_fields: NoParams =
3020 serde_json::from_str(r#"{"unexpected": "value"}"#).unwrap();
3021 assert_eq!(from_object_with_fields, NoParams);
3022 }
3023
3024 #[tokio::test]
3025 async fn test_no_params_type_in_handler() {
3026 let tool = ToolBuilder::new("status")
3028 .description("Get status")
3029 .handler(|_input: NoParams| async move { Ok(CallToolResult::text("OK")) })
3030 .build();
3031
3032 let schema = tool.definition().input_schema;
3034 assert_eq!(
3035 schema.get("type").and_then(|v| v.as_str()),
3036 Some("object"),
3037 "NoParams handler should produce type: object schema"
3038 );
3039
3040 let result = tool.call(serde_json::json!({})).await;
3042 assert!(!result.is_error);
3043 }
3044
3045 #[tokio::test]
3046 async fn test_serde_json_value_handler_has_type_object() {
3047 let tool = ToolBuilder::new("any_input")
3050 .description("Accepts any input")
3051 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3052 .build();
3053
3054 let schema = tool.definition().input_schema;
3055 assert_eq!(
3056 schema.get("type").and_then(|v| v.as_str()),
3057 Some("object"),
3058 "serde_json::Value handler should produce schema with type: object"
3059 );
3060 }
3061
3062 #[tokio::test]
3063 async fn test_tool_with_name_prefix() {
3064 #[derive(Debug, Deserialize, JsonSchema)]
3065 struct Input {
3066 value: String,
3067 }
3068
3069 let tool = ToolBuilder::new("query")
3070 .description("Query something")
3071 .title("Query Tool")
3072 .handler(|input: Input| async move { Ok(CallToolResult::text(&input.value)) })
3073 .build();
3074
3075 let prefixed = tool.with_name_prefix("db");
3077
3078 assert_eq!(prefixed.name, "db.query");
3080
3081 assert_eq!(prefixed.description.as_deref(), Some("Query something"));
3083 assert_eq!(prefixed.title.as_deref(), Some("Query Tool"));
3084
3085 let result = prefixed
3087 .call(serde_json::json!({"value": "test input"}))
3088 .await;
3089 assert!(!result.is_error);
3090 match &result.content[0] {
3091 Content::Text { text, .. } => assert_eq!(text, "test input"),
3092 _ => panic!("Expected text content"),
3093 }
3094 }
3095
3096 #[tokio::test]
3097 async fn test_tool_with_name_prefix_multiple_levels() {
3098 let tool = ToolBuilder::new("action")
3099 .description("Do something")
3100 .handler(|_: NoParams| async move { Ok(CallToolResult::text("done")) })
3101 .build();
3102
3103 let prefixed = tool.with_name_prefix("level1");
3105 assert_eq!(prefixed.name, "level1.action");
3106
3107 let double_prefixed = prefixed.with_name_prefix("level0");
3108 assert_eq!(double_prefixed.name, "level0.level1.action");
3109 }
3110
3111 #[tokio::test]
3116 async fn test_no_params_handler_basic() {
3117 let tool = ToolBuilder::new("get_status")
3118 .description("Get current status")
3119 .no_params_handler(|| async { Ok(CallToolResult::text("OK")) })
3120 .build();
3121
3122 assert_eq!(tool.name, "get_status");
3123 assert_eq!(tool.description.as_deref(), Some("Get current status"));
3124
3125 let result = tool.call(serde_json::json!({})).await;
3127 assert!(!result.is_error);
3128 assert_eq!(result.first_text().unwrap(), "OK");
3129
3130 let result = tool.call(serde_json::json!(null)).await;
3132 assert!(!result.is_error);
3133
3134 let schema = tool.definition().input_schema;
3136 assert_eq!(schema.get("type").and_then(|v| v.as_str()), Some("object"));
3137 }
3138
3139 #[tokio::test]
3140 async fn test_no_params_handler_with_captured_state() {
3141 let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
3142 let counter_ref = counter.clone();
3143
3144 let tool = ToolBuilder::new("increment")
3145 .description("Increment counter")
3146 .no_params_handler(move || {
3147 let c = counter_ref.clone();
3148 async move {
3149 let prev = c.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
3150 Ok(CallToolResult::text(format!("Incremented from {}", prev)))
3151 }
3152 })
3153 .build();
3154
3155 let _ = tool.call(serde_json::json!({})).await;
3157 let _ = tool.call(serde_json::json!({})).await;
3158 let result = tool.call(serde_json::json!({})).await;
3159
3160 assert!(!result.is_error);
3161 assert_eq!(result.first_text().unwrap(), "Incremented from 2");
3162 assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 3);
3163 }
3164
3165 #[tokio::test]
3166 async fn test_no_params_handler_with_layer() {
3167 use std::time::Duration;
3168 use tower::timeout::TimeoutLayer;
3169
3170 let tool = ToolBuilder::new("slow_status")
3171 .description("Slow status check")
3172 .no_params_handler(|| async {
3173 tokio::time::sleep(Duration::from_millis(10)).await;
3174 Ok(CallToolResult::text("done"))
3175 })
3176 .layer(TimeoutLayer::new(Duration::from_secs(1)))
3177 .build();
3178
3179 let result = tool.call(serde_json::json!({})).await;
3180 assert!(!result.is_error);
3181 assert_eq!(result.first_text().unwrap(), "done");
3182 }
3183
3184 #[tokio::test]
3185 async fn test_no_params_handler_timeout() {
3186 use std::time::Duration;
3187 use tower::timeout::TimeoutLayer;
3188
3189 let tool = ToolBuilder::new("very_slow_status")
3190 .description("Very slow status check")
3191 .no_params_handler(|| async {
3192 tokio::time::sleep(Duration::from_millis(200)).await;
3193 Ok(CallToolResult::text("done"))
3194 })
3195 .layer(TimeoutLayer::new(Duration::from_millis(50)))
3196 .build();
3197
3198 let result = tool.call(serde_json::json!({})).await;
3199 assert!(result.is_error);
3200 let msg = result.first_text().unwrap().to_lowercase();
3201 assert!(
3202 msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
3203 "Expected timeout error, got: {}",
3204 msg
3205 );
3206 }
3207
3208 #[tokio::test]
3209 async fn test_no_params_handler_with_multiple_layers() {
3210 use std::time::Duration;
3211 use tower::limit::ConcurrencyLimitLayer;
3212 use tower::timeout::TimeoutLayer;
3213
3214 let tool = ToolBuilder::new("multi_layer_status")
3215 .description("Status with multiple layers")
3216 .no_params_handler(|| async { Ok(CallToolResult::text("status ok")) })
3217 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3218 .layer(ConcurrencyLimitLayer::new(10))
3219 .build();
3220
3221 let result = tool.call(serde_json::json!({})).await;
3222 assert!(!result.is_error);
3223 assert_eq!(result.first_text().unwrap(), "status ok");
3224 }
3225
3226 #[tokio::test]
3231 async fn test_guard_allows_request() {
3232 #[derive(Debug, Deserialize, JsonSchema)]
3233 #[allow(dead_code)]
3234 struct DeleteInput {
3235 id: String,
3236 confirm: bool,
3237 }
3238
3239 let tool = ToolBuilder::new("delete")
3240 .description("Delete a record")
3241 .handler(|input: DeleteInput| async move {
3242 Ok(CallToolResult::text(format!("deleted {}", input.id)))
3243 })
3244 .guard(|req: &ToolRequest| {
3245 let confirm = req
3246 .args
3247 .get("confirm")
3248 .and_then(|v| v.as_bool())
3249 .unwrap_or(false);
3250 if !confirm {
3251 return Err("Must set confirm=true to delete".to_string());
3252 }
3253 Ok(())
3254 })
3255 .build();
3256
3257 let result = tool
3258 .call(serde_json::json!({"id": "abc", "confirm": true}))
3259 .await;
3260 assert!(!result.is_error);
3261 assert_eq!(result.first_text().unwrap(), "deleted abc");
3262 }
3263
3264 #[tokio::test]
3265 async fn test_guard_rejects_request() {
3266 #[derive(Debug, Deserialize, JsonSchema)]
3267 #[allow(dead_code)]
3268 struct DeleteInput2 {
3269 id: String,
3270 confirm: bool,
3271 }
3272
3273 let tool = ToolBuilder::new("delete2")
3274 .description("Delete a record")
3275 .handler(|input: DeleteInput2| async move {
3276 Ok(CallToolResult::text(format!("deleted {}", input.id)))
3277 })
3278 .guard(|req: &ToolRequest| {
3279 let confirm = req
3280 .args
3281 .get("confirm")
3282 .and_then(|v| v.as_bool())
3283 .unwrap_or(false);
3284 if !confirm {
3285 return Err("Must set confirm=true to delete".to_string());
3286 }
3287 Ok(())
3288 })
3289 .build();
3290
3291 let result = tool
3292 .call(serde_json::json!({"id": "abc", "confirm": false}))
3293 .await;
3294 assert!(result.is_error);
3295 assert!(
3296 result
3297 .first_text()
3298 .unwrap()
3299 .contains("Must set confirm=true")
3300 );
3301 }
3302
3303 #[tokio::test]
3304 async fn test_guard_with_layer() {
3305 use std::time::Duration;
3306 use tower::timeout::TimeoutLayer;
3307
3308 let tool = ToolBuilder::new("guarded_timeout")
3309 .description("Guarded with timeout")
3310 .handler(|input: GreetInput| async move {
3311 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3312 })
3313 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3314 .guard(|_req: &ToolRequest| Ok(()))
3315 .build();
3316
3317 let result = tool.call(serde_json::json!({"name": "World"})).await;
3318 assert!(!result.is_error);
3319 assert_eq!(result.first_text().unwrap(), "Hello, World!");
3320 }
3321
3322 #[tokio::test]
3323 async fn test_guard_on_no_params_handler() {
3324 let allowed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
3325 let allowed_clone = allowed.clone();
3326
3327 let tool = ToolBuilder::new("status")
3328 .description("Get status")
3329 .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3330 .guard(move |_req: &ToolRequest| {
3331 if allowed_clone.load(std::sync::atomic::Ordering::Relaxed) {
3332 Ok(())
3333 } else {
3334 Err("Access denied".to_string())
3335 }
3336 })
3337 .build();
3338
3339 let result = tool.call(serde_json::json!({})).await;
3341 assert!(!result.is_error);
3342 assert_eq!(result.first_text().unwrap(), "ok");
3343
3344 allowed.store(false, std::sync::atomic::Ordering::Relaxed);
3346 let result = tool.call(serde_json::json!({})).await;
3347 assert!(result.is_error);
3348 assert!(result.first_text().unwrap().contains("Access denied"));
3349 }
3350
3351 #[tokio::test]
3352 async fn test_guard_on_no_params_handler_with_layer() {
3353 use std::time::Duration;
3354 use tower::timeout::TimeoutLayer;
3355
3356 let tool = ToolBuilder::new("status_layered")
3357 .description("Get status with layers")
3358 .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3359 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3360 .guard(|_req: &ToolRequest| Ok(()))
3361 .build();
3362
3363 let result = tool.call(serde_json::json!({})).await;
3364 assert!(!result.is_error);
3365 assert_eq!(result.first_text().unwrap(), "ok");
3366 }
3367
3368 #[tokio::test]
3369 async fn test_guard_on_extractor_handler() {
3370 use std::sync::Arc;
3371
3372 #[derive(Clone)]
3373 struct AppState {
3374 prefix: String,
3375 }
3376
3377 #[derive(Debug, Deserialize, JsonSchema)]
3378 struct QueryInput {
3379 query: String,
3380 }
3381
3382 let state = Arc::new(AppState {
3383 prefix: "db".to_string(),
3384 });
3385
3386 let tool = ToolBuilder::new("search")
3387 .description("Search")
3388 .extractor_handler(
3389 state,
3390 |State(app): State<Arc<AppState>>, Json(input): Json<QueryInput>| async move {
3391 Ok(CallToolResult::text(format!(
3392 "{}: {}",
3393 app.prefix, input.query
3394 )))
3395 },
3396 )
3397 .guard(|req: &ToolRequest| {
3398 let query = req.args.get("query").and_then(|v| v.as_str()).unwrap_or("");
3399 if query.is_empty() {
3400 return Err("Query cannot be empty".to_string());
3401 }
3402 Ok(())
3403 })
3404 .build();
3405
3406 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3408 assert!(!result.is_error);
3409 assert_eq!(result.first_text().unwrap(), "db: hello");
3410
3411 let result = tool.call(serde_json::json!({"query": ""})).await;
3413 assert!(result.is_error);
3414 assert!(
3415 result
3416 .first_text()
3417 .unwrap()
3418 .contains("Query cannot be empty")
3419 );
3420 }
3421
3422 #[tokio::test]
3423 async fn test_guard_on_extractor_handler_with_layer() {
3424 use std::sync::Arc;
3425 use std::time::Duration;
3426 use tower::timeout::TimeoutLayer;
3427
3428 #[derive(Clone)]
3429 struct AppState2 {
3430 prefix: String,
3431 }
3432
3433 #[derive(Debug, Deserialize, JsonSchema)]
3434 struct QueryInput2 {
3435 query: String,
3436 }
3437
3438 let state = Arc::new(AppState2 {
3439 prefix: "db".to_string(),
3440 });
3441
3442 let tool = ToolBuilder::new("search2")
3443 .description("Search with layer and guard")
3444 .extractor_handler(
3445 state,
3446 |State(app): State<Arc<AppState2>>, Json(input): Json<QueryInput2>| async move {
3447 Ok(CallToolResult::text(format!(
3448 "{}: {}",
3449 app.prefix, input.query
3450 )))
3451 },
3452 )
3453 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3454 .guard(|_req: &ToolRequest| Ok(()))
3455 .build();
3456
3457 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3458 assert!(!result.is_error);
3459 assert_eq!(result.first_text().unwrap(), "db: hello");
3460 }
3461
3462 #[tokio::test]
3463 async fn test_tool_with_guard_post_build() {
3464 let tool = ToolBuilder::new("admin_action")
3465 .description("Admin action")
3466 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3467 .build();
3468
3469 let guarded = tool.with_guard(|req: &ToolRequest| {
3471 let name = req.args.get("name").and_then(|v| v.as_str()).unwrap_or("");
3472 if name == "admin" {
3473 Ok(())
3474 } else {
3475 Err("Only admin allowed".to_string())
3476 }
3477 });
3478
3479 let result = guarded.call(serde_json::json!({"name": "admin"})).await;
3481 assert!(!result.is_error);
3482
3483 let result = guarded.call(serde_json::json!({"name": "user"})).await;
3485 assert!(result.is_error);
3486 assert!(result.first_text().unwrap().contains("Only admin allowed"));
3487 }
3488
3489 #[tokio::test]
3490 async fn test_with_guard_preserves_tool_metadata() {
3491 let tool = ToolBuilder::new("my_tool")
3492 .description("A tool")
3493 .title("My Tool")
3494 .read_only()
3495 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3496 .build();
3497
3498 let guarded = tool.with_guard(|_req: &ToolRequest| Ok(()));
3499
3500 assert_eq!(guarded.name, "my_tool");
3501 assert_eq!(guarded.description.as_deref(), Some("A tool"));
3502 assert_eq!(guarded.title.as_deref(), Some("My Tool"));
3503 assert!(guarded.annotations.is_some());
3504 }
3505
3506 #[tokio::test]
3507 async fn test_guard_group_pattern() {
3508 let require_auth = |req: &ToolRequest| {
3510 let token = req
3511 .args
3512 .get("_token")
3513 .and_then(|v| v.as_str())
3514 .unwrap_or("");
3515 if token == "valid" {
3516 Ok(())
3517 } else {
3518 Err("Authentication required".to_string())
3519 }
3520 };
3521
3522 let tool1 = ToolBuilder::new("action1")
3523 .description("Action 1")
3524 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action1")) })
3525 .build();
3526 let tool2 = ToolBuilder::new("action2")
3527 .description("Action 2")
3528 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action2")) })
3529 .build();
3530
3531 let guarded1 = tool1.with_guard(require_auth);
3533 let guarded2 = tool2.with_guard(require_auth);
3534
3535 let r1 = guarded1
3537 .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3538 .await;
3539 let r2 = guarded2
3540 .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3541 .await;
3542 assert!(r1.is_error);
3543 assert!(r2.is_error);
3544
3545 let r1 = guarded1
3547 .call(serde_json::json!({"name": "test", "_token": "valid"}))
3548 .await;
3549 let r2 = guarded2
3550 .call(serde_json::json!({"name": "test", "_token": "valid"}))
3551 .await;
3552 assert!(!r1.is_error);
3553 assert!(!r2.is_error);
3554 }
3555
3556 #[tokio::test]
3557 async fn test_input_validation_returns_tool_error() {
3558 #[derive(Debug, Deserialize, JsonSchema)]
3561 struct StrictInput {
3562 name: String,
3563 count: u32,
3564 }
3565
3566 let tool = ToolBuilder::new("strict_tool")
3567 .description("requires specific input")
3568 .handler(|input: StrictInput| async move {
3569 Ok(CallToolResult::text(format!(
3570 "{}: {}",
3571 input.name, input.count
3572 )))
3573 })
3574 .build();
3575
3576 let result = tool
3578 .call(serde_json::json!({"name": "test", "count": 5}))
3579 .await;
3580 assert!(!result.is_error);
3581
3582 let result = tool.call(serde_json::json!({"name": "test"})).await;
3584 assert!(result.is_error);
3585 let text = result.first_text().unwrap();
3586 assert!(text.contains("Invalid input"), "got: {text}");
3587
3588 let result = tool
3590 .call(serde_json::json!({"name": "test", "count": "not_a_number"}))
3591 .await;
3592 assert!(result.is_error);
3593 let text = result.first_text().unwrap();
3594 assert!(text.contains("Invalid input"), "got: {text}");
3595 }
3596
3597 #[tokio::test]
3598 async fn test_input_schema_override_with_raw_args() {
3599 let custom = serde_json::json!({
3603 "type": "object",
3604 "properties": {
3605 "query": { "type": "string", "minLength": 1 }
3606 },
3607 "required": ["query"]
3608 });
3609
3610 let tool = ToolBuilder::new("query")
3611 .description("Query with a custom schema")
3612 .input_schema(custom.clone())
3613 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3614 Ok(CallToolResult::json(args))
3615 })
3616 .build();
3617
3618 let schema = tool.definition().input_schema;
3619 assert_eq!(schema, custom);
3620
3621 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3623 assert!(!result.is_error);
3624 }
3625
3626 #[tokio::test]
3627 async fn test_input_schema_override_wins_over_typed_handler() {
3628 let custom = serde_json::json!({
3632 "type": "object",
3633 "title": "GreetOverride",
3634 "properties": {
3635 "name": { "type": "string", "minLength": 1, "maxLength": 64 }
3636 },
3637 "required": ["name"],
3638 "additionalProperties": false
3639 });
3640
3641 let tool = ToolBuilder::new("greet")
3642 .description("Greet someone with a hand-tuned schema")
3643 .input_schema(custom.clone())
3644 .handler(|input: GreetInput| async move {
3645 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3646 })
3647 .build();
3648
3649 let schema = tool.definition().input_schema;
3650 assert_eq!(schema, custom);
3651 assert_eq!(schema["title"], "GreetOverride");
3653
3654 let result = tool.call(serde_json::json!({"name": "World"})).await;
3656 assert!(!result.is_error);
3657 }
3658
3659 #[tokio::test]
3660 async fn test_input_schema_override_preserves_2020_12_constructs() {
3661 let custom = serde_json::json!({
3664 "type": "object",
3665 "properties": {
3666 "filter": {
3667 "oneOf": [
3668 { "type": "string" },
3669 {
3670 "type": "object",
3671 "properties": { "field": { "type": "string" } },
3672 "required": ["field"]
3673 }
3674 ]
3675 }
3676 },
3677 "required": ["filter"]
3678 });
3679
3680 let tool = ToolBuilder::new("filter_tool")
3681 .description("Demonstrates oneOf preservation")
3682 .input_schema(custom.clone())
3683 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3684 Ok(CallToolResult::json(args))
3685 })
3686 .build();
3687
3688 let schema = tool.definition().input_schema;
3689 assert_eq!(schema, custom);
3690 let one_of = schema["properties"]["filter"]["oneOf"]
3691 .as_array()
3692 .expect("oneOf must survive as an array");
3693 assert_eq!(one_of.len(), 2);
3694 assert_eq!(one_of[0]["type"], "string");
3695 assert_eq!(one_of[1]["type"], "object");
3696 }
3697
3698 #[tokio::test]
3699 async fn test_input_schema_override_adds_type_object_if_missing() {
3700 let custom_no_type = serde_json::json!({
3703 "properties": {
3704 "x": { "type": "number" }
3705 }
3706 });
3707
3708 let tool = ToolBuilder::new("typeless")
3709 .description("Schema missing top-level type")
3710 .input_schema(custom_no_type)
3711 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3712 Ok(CallToolResult::json(args))
3713 })
3714 .build();
3715
3716 let schema = tool.definition().input_schema;
3717 assert_eq!(schema["type"], "object");
3718 assert!(schema["properties"]["x"].is_object());
3719 }
3720}