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::{Map, 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::{Extensions, 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
461#[derive(Debug, Clone, PartialEq, Eq)]
466pub struct TaskContext {
467 task_id: String,
468}
469
470impl TaskContext {
471 pub(crate) fn new(task_id: String) -> Self {
472 Self { task_id }
473 }
474
475 pub fn task_id(&self) -> &str {
477 &self.task_id
478 }
479}
480
481#[derive(Debug, Clone, Default)]
483pub struct TaskPreparation {
484 pub(crate) meta: Option<Map<String, Value>>,
485 pub(crate) extensions: Extensions,
486}
487
488impl TaskPreparation {
489 pub fn new() -> Self {
491 Self::default()
492 }
493
494 pub fn with_meta(mut self, meta: Map<String, Value>) -> Self {
496 self.meta = Some(meta);
497 self
498 }
499
500 pub fn with_extension<T: Send + Sync + 'static>(mut self, value: T) -> Self {
503 self.extensions.insert(value);
504 self
505 }
506}
507
508pub(crate) trait TaskPreparer: Send + Sync {
509 fn prepare(
510 &self,
511 context: TaskContext,
512 arguments: Value,
513 ) -> BoxFuture<'_, Result<TaskPreparation>>;
514}
515
516impl<F, Fut> TaskPreparer for F
517where
518 F: Fn(TaskContext, Value) -> Fut + Send + Sync,
519 Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
520{
521 fn prepare(
522 &self,
523 context: TaskContext,
524 arguments: Value,
525 ) -> BoxFuture<'_, Result<TaskPreparation>> {
526 Box::pin((self)(context, arguments))
527 }
528}
529
530struct TypedTaskPreparer<I, F> {
531 prepare: F,
532 _phantom: std::marker::PhantomData<I>,
533}
534
535impl<I, F, Fut> TaskPreparer for TypedTaskPreparer<I, F>
536where
537 I: DeserializeOwned + Send + Sync + 'static,
538 F: Fn(TaskContext, I) -> Fut + Send + Sync,
539 Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
540{
541 fn prepare(
542 &self,
543 context: TaskContext,
544 arguments: Value,
545 ) -> BoxFuture<'_, Result<TaskPreparation>> {
546 let input = serde_json::from_value(arguments)
547 .map_err(|error| Error::invalid_params(format!("Invalid input: {error}")));
548 match input {
549 Ok(input) => Box::pin((self.prepare)(context, input)),
550 Err(error) => Box::pin(async move { Err(error) }),
551 }
552 }
553}
554
555pub trait ToolHandler: Send + Sync {
557 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>>;
559
560 fn call_with_context(
565 &self,
566 _ctx: RequestContext,
567 args: Value,
568 ) -> BoxFuture<'_, Result<CallToolResult>> {
569 self.call(args)
570 }
571
572 fn uses_context(&self) -> bool {
574 false
575 }
576
577 fn input_schema(&self) -> Value;
579}
580
581#[cfg(feature = "stateless")]
584pub trait MrtrToolHandler: Send + Sync {
585 fn call(
587 &self,
588 ctx: RequestContext,
589 args: Value,
590 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>>;
591
592 fn input_schema(&self) -> Value;
594}
595
596#[cfg(feature = "stateless")]
598struct MrtrToolHandlerService<H> {
599 handler: Arc<H>,
600}
601
602#[cfg(feature = "stateless")]
603impl<H> MrtrToolHandlerService<H> {
604 fn new(handler: H) -> Self {
605 Self {
606 handler: Arc::new(handler),
607 }
608 }
609}
610
611#[cfg(feature = "stateless")]
612impl<H> Clone for MrtrToolHandlerService<H> {
613 fn clone(&self) -> Self {
614 Self {
615 handler: self.handler.clone(),
616 }
617 }
618}
619
620#[cfg(feature = "stateless")]
621impl<H> Service<ToolRequest> for MrtrToolHandlerService<H>
622where
623 H: MrtrToolHandler + 'static,
624{
625 type Response = RequestOutcome<CallToolResult>;
626 type Error = Error;
627 type Future =
628 Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
629
630 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
631 Poll::Ready(Ok(()))
632 }
633
634 fn call(&mut self, req: ToolRequest) -> Self::Future {
635 let handler = self.handler.clone();
636 Box::pin(async move { handler.call(req.ctx, req.args).await })
637 }
638}
639
640#[cfg(feature = "stateless")]
642struct ServiceMrtrToolHandler {
643 service: Mutex<BoxMrtrToolService>,
644 input_schema: Value,
645}
646
647#[cfg(feature = "stateless")]
648struct GuardedMrtrToolHandler<G> {
649 guard: G,
650 inner: Arc<dyn MrtrToolHandler>,
651}
652
653#[cfg(feature = "stateless")]
654impl<G> MrtrToolHandler for GuardedMrtrToolHandler<G>
655where
656 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
657{
658 fn call(
659 &self,
660 ctx: RequestContext,
661 args: Value,
662 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
663 let request = ToolRequest::new(ctx, args);
664 match (self.guard)(&request) {
665 Ok(()) => self.inner.call(request.ctx, request.args),
666 Err(message) => {
667 Box::pin(
668 async move { Ok(RequestOutcome::Complete(CallToolResult::error(message))) },
669 )
670 }
671 }
672 }
673
674 fn input_schema(&self) -> Value {
675 self.inner.input_schema()
676 }
677}
678
679#[cfg(feature = "stateless")]
680impl MrtrToolHandler for ServiceMrtrToolHandler {
681 fn call(
682 &self,
683 ctx: RequestContext,
684 args: Value,
685 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
686 Box::pin(async move {
687 let mut service = self.service.lock().await.clone();
688 let outcome = service
689 .ready()
690 .await
691 .expect("MRTR tool service is infallible")
692 .call(ToolRequest::new(ctx, args))
693 .await
694 .expect("MRTR tool service is infallible");
695 Ok(outcome)
696 })
697 }
698
699 fn input_schema(&self) -> Value {
700 self.input_schema.clone()
701 }
702}
703
704pub(crate) struct ToolHandlerService<H> {
709 handler: Arc<H>,
710}
711
712impl<H> ToolHandlerService<H> {
713 pub(crate) fn new(handler: H) -> Self {
714 Self {
715 handler: Arc::new(handler),
716 }
717 }
718}
719
720impl<H> Clone for ToolHandlerService<H> {
721 fn clone(&self) -> Self {
722 Self {
723 handler: self.handler.clone(),
724 }
725 }
726}
727
728impl<H> Service<ToolRequest> for ToolHandlerService<H>
729where
730 H: ToolHandler + 'static,
731{
732 type Response = CallToolResult;
733 type Error = Error;
734 type Future = Pin<Box<dyn Future<Output = std::result::Result<CallToolResult, Error>> + Send>>;
735
736 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
737 Poll::Ready(Ok(()))
738 }
739
740 fn call(&mut self, req: ToolRequest) -> Self::Future {
741 let handler = self.handler.clone();
742 Box::pin(async move { handler.call_with_context(req.ctx, req.args).await })
743 }
744}
745
746pub struct Tool {
753 pub name: String,
755 pub title: Option<String>,
757 pub description: Option<String>,
759 pub output_schema: Option<Value>,
761 pub icons: Option<Vec<ToolIcon>>,
763 pub annotations: Option<ToolAnnotations>,
765 pub meta: Option<Value>,
767 pub task_support: TaskSupportMode,
769 pub(crate) required_client_capabilities: Option<ClientCapabilities>,
772 pub(crate) task_preparer: Option<Arc<dyn TaskPreparer>>,
774 pub(crate) service: Option<BoxToolService>,
776 #[cfg(feature = "stateless")]
777 pub(crate) mrtr_handler: Option<Arc<dyn MrtrToolHandler>>,
778 pub(crate) input_schema: Value,
780}
781
782impl std::fmt::Debug for Tool {
783 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
784 f.debug_struct("Tool")
785 .field("name", &self.name)
786 .field("title", &self.title)
787 .field("description", &self.description)
788 .field("output_schema", &self.output_schema)
789 .field("icons", &self.icons)
790 .field("annotations", &self.annotations)
791 .field("meta", &self.meta)
792 .field("task_support", &self.task_support)
793 .field(
794 "required_client_capabilities",
795 &self.required_client_capabilities,
796 )
797 .finish_non_exhaustive()
798 }
799}
800
801unsafe impl Send for Tool {}
804unsafe impl Sync for Tool {}
805
806impl Clone for Tool {
807 fn clone(&self) -> Self {
808 Self {
809 name: self.name.clone(),
810 title: self.title.clone(),
811 description: self.description.clone(),
812 output_schema: self.output_schema.clone(),
813 icons: self.icons.clone(),
814 annotations: self.annotations.clone(),
815 meta: self.meta.clone(),
816 task_support: self.task_support,
817 required_client_capabilities: self.required_client_capabilities.clone(),
818 task_preparer: self.task_preparer.clone(),
819 service: self.service.clone(),
820 #[cfg(feature = "stateless")]
821 mrtr_handler: self.mrtr_handler.clone(),
822 input_schema: self.input_schema.clone(),
823 }
824 }
825}
826
827impl Tool {
828 pub fn builder(name: impl Into<String>) -> ToolBuilder {
830 ToolBuilder::new(name)
831 }
832
833 pub fn definition(&self) -> ToolDefinition {
835 let execution = match self.task_support {
836 TaskSupportMode::Forbidden => None,
837 mode => Some(ToolExecution {
838 task_support: Some(mode),
839 }),
840 };
841 ToolDefinition {
842 name: self.name.clone(),
843 title: self.title.clone(),
844 description: self.description.clone(),
845 input_schema: self.input_schema.clone(),
846 output_schema: self.output_schema.clone(),
847 icons: self.icons.clone(),
848 annotations: self.annotations.clone(),
849 execution,
850 meta: self.meta.clone(),
851 }
852 }
853
854 pub fn with_meta(
856 mut self,
857 meta: Value,
858 ) -> std::result::Result<Self, crate::protocol::MetaValidationError> {
859 crate::protocol::validate_meta_object(&meta)?;
860 self.meta = Some(meta);
861 Ok(self)
862 }
863
864 pub fn call(&self, args: Value) -> BoxFuture<'static, CallToolResult> {
869 let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
870 self.call_with_context(ctx, args)
871 }
872
873 pub fn call_with_context(
884 &self,
885 ctx: RequestContext,
886 args: Value,
887 ) -> BoxFuture<'static, CallToolResult> {
888 let tool = self.clone();
889 Box::pin(async move {
890 match tool.call_outcome_with_context(ctx, args).await {
891 Ok(RequestOutcome::Complete(result)) => result,
892 Ok(RequestOutcome::InputRequired(_)) => CallToolResult::error(
893 "tool requires additional client input; use call_outcome_with_context",
894 ),
895 Err(error) => CallToolResult::error(error.to_string()),
896 }
897 })
898 }
899
900 pub fn call_outcome(
902 &self,
903 args: Value,
904 ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
905 let ctx = RequestContext::new(crate::protocol::RequestId::Number(0));
906 self.call_outcome_with_context(ctx, args)
907 }
908
909 pub fn call_outcome_with_context(
912 &self,
913 ctx: RequestContext,
914 args: Value,
915 ) -> BoxFuture<'static, Result<RequestOutcome<CallToolResult>>> {
916 use tower::ServiceExt;
917 #[cfg(feature = "stateless")]
918 if let Some(handler) = self.mrtr_handler.clone() {
919 return Box::pin(async move { handler.call(ctx, args).await });
920 }
921 let service = self
922 .service
923 .clone()
924 .expect("tool must have a complete or MRTR handler");
925 Box::pin(async move {
926 let result = service.oneshot(ToolRequest::new(ctx, args)).await.unwrap();
927 Ok(RequestOutcome::Complete(result))
928 })
929 }
930
931 pub fn require_client_capabilities(mut self, required: ClientCapabilities) -> Self {
939 self.required_client_capabilities = Some(required);
940 self
941 }
942
943 pub fn required_client_capabilities(&self) -> Option<&ClientCapabilities> {
945 self.required_client_capabilities.as_ref()
946 }
947
948 pub fn with_task_preparation<F, Fut>(mut self, prepare: F) -> Self
953 where
954 F: Fn(TaskContext, Value) -> Fut + Send + Sync + 'static,
955 Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
956 {
957 self.task_preparer = Some(Arc::new(prepare));
958 self
959 }
960
961 pub fn with_typed_task_preparation<I, F, Fut>(mut self, prepare: F) -> Self
963 where
964 I: DeserializeOwned + Send + Sync + 'static,
965 F: Fn(TaskContext, I) -> Fut + Send + Sync + 'static,
966 Fut: Future<Output = Result<TaskPreparation>> + Send + 'static,
967 {
968 self.task_preparer = Some(Arc::new(TypedTaskPreparer {
969 prepare,
970 _phantom: std::marker::PhantomData,
971 }));
972 self
973 }
974
975 pub(crate) async fn prepare_task(
976 &self,
977 context: TaskContext,
978 arguments: Value,
979 ) -> Result<TaskPreparation> {
980 match self.task_preparer.as_ref() {
981 Some(prepare) => prepare.prepare(context, arguments).await,
982 None => Ok(TaskPreparation::default()),
983 }
984 }
985
986 pub fn with_guard<G>(self, guard: G) -> Self
1015 where
1016 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1017 {
1018 #[cfg(feature = "stateless")]
1019 if let Some(inner) = self.mrtr_handler.clone() {
1020 return Tool {
1021 mrtr_handler: Some(Arc::new(GuardedMrtrToolHandler { guard, inner })),
1022 ..self
1023 };
1024 }
1025
1026 let guarded = GuardService {
1027 guard,
1028 inner: self
1029 .service
1030 .expect("tool must have a complete or MRTR handler"),
1031 };
1032 let caught = ToolCatchError::new(guarded);
1033 Tool {
1034 service: Some(BoxCloneService::new(caught)),
1035 ..self
1036 }
1037 }
1038
1039 pub fn with_name_prefix(&self, prefix: &str) -> Self {
1066 Self {
1067 name: format!("{}.{}", prefix, self.name),
1068 title: self.title.clone(),
1069 description: self.description.clone(),
1070 output_schema: self.output_schema.clone(),
1071 icons: self.icons.clone(),
1072 annotations: self.annotations.clone(),
1073 meta: self.meta.clone(),
1074 task_support: self.task_support,
1075 required_client_capabilities: self.required_client_capabilities.clone(),
1076 task_preparer: self.task_preparer.clone(),
1077 service: self.service.clone(),
1078 #[cfg(feature = "stateless")]
1079 mrtr_handler: self.mrtr_handler.clone(),
1080 input_schema: self.input_schema.clone(),
1081 }
1082 }
1083
1084 #[allow(clippy::too_many_arguments)]
1086 fn from_handler<H: ToolHandler + 'static>(
1087 name: String,
1088 title: Option<String>,
1089 description: Option<String>,
1090 output_schema: Option<Value>,
1091 icons: Option<Vec<ToolIcon>>,
1092 annotations: Option<ToolAnnotations>,
1093 task_support: TaskSupportMode,
1094 input_schema_override: Option<Value>,
1095 handler: H,
1096 ) -> Self {
1097 let input_schema =
1098 ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
1099 let handler_service = ToolHandlerService::new(handler);
1100 let catch_error = ToolCatchError::new(handler_service);
1101 let service = BoxCloneService::new(catch_error);
1102
1103 Self {
1104 name,
1105 title,
1106 description,
1107 output_schema,
1108 icons,
1109 annotations,
1110 meta: None,
1111 task_support,
1112 required_client_capabilities: None,
1113 task_preparer: None,
1114 service: Some(service),
1115 #[cfg(feature = "stateless")]
1116 mrtr_handler: None,
1117 input_schema,
1118 }
1119 }
1120
1121 #[cfg(feature = "stateless")]
1122 #[allow(clippy::too_many_arguments)]
1123 fn from_mrtr_handler<H: MrtrToolHandler + 'static>(
1124 name: String,
1125 title: Option<String>,
1126 description: Option<String>,
1127 output_schema: Option<Value>,
1128 icons: Option<Vec<ToolIcon>>,
1129 annotations: Option<ToolAnnotations>,
1130 task_support: TaskSupportMode,
1131 input_schema_override: Option<Value>,
1132 handler: H,
1133 ) -> Self {
1134 let input_schema =
1135 ensure_object_schema(input_schema_override.unwrap_or_else(|| handler.input_schema()));
1136 Self {
1137 name,
1138 title,
1139 description,
1140 output_schema,
1141 icons,
1142 annotations,
1143 meta: None,
1144 task_support,
1145 required_client_capabilities: None,
1146 task_preparer: None,
1147 service: None,
1148 mrtr_handler: Some(Arc::new(handler)),
1149 input_schema,
1150 }
1151 }
1152}
1153
1154pub struct ToolBuilder {
1182 name: String,
1183 title: Option<String>,
1184 description: Option<String>,
1185 output_schema: Option<Value>,
1186 input_schema_override: Option<Value>,
1187 icons: Option<Vec<ToolIcon>>,
1188 annotations: Option<ToolAnnotations>,
1189 task_support: TaskSupportMode,
1190}
1191
1192impl ToolBuilder {
1193 pub fn new(name: impl Into<String>) -> Self {
1206 let name = name.into();
1207 if let Err(e) = validate_tool_name(&name) {
1208 panic!("{e}");
1209 }
1210 Self {
1211 name,
1212 title: None,
1213 description: None,
1214 output_schema: None,
1215 input_schema_override: None,
1216 icons: None,
1217 annotations: None,
1218 task_support: TaskSupportMode::default(),
1219 }
1220 }
1221
1222 pub fn try_new(name: impl Into<String>) -> Result<Self> {
1228 let name = name.into();
1229 validate_tool_name(&name)?;
1230 Ok(Self {
1231 name,
1232 title: None,
1233 description: None,
1234 output_schema: None,
1235 input_schema_override: None,
1236 icons: None,
1237 annotations: None,
1238 task_support: TaskSupportMode::default(),
1239 })
1240 }
1241
1242 pub fn title(mut self, title: impl Into<String>) -> Self {
1258 self.title = Some(title.into());
1259 self
1260 }
1261
1262 pub fn output_schema(mut self, schema: Value) -> Self {
1264 self.output_schema = Some(schema);
1265 self
1266 }
1267
1268 pub fn input_schema(mut self, schema: Value) -> Self {
1321 self.input_schema_override = Some(schema);
1322 self
1323 }
1324
1325 pub fn icon(mut self, src: impl Into<String>) -> Self {
1327 self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1328 src: src.into(),
1329 mime_type: None,
1330 sizes: None,
1331 theme: None,
1332 });
1333 self
1334 }
1335
1336 pub fn icon_with_meta(
1338 mut self,
1339 src: impl Into<String>,
1340 mime_type: Option<String>,
1341 sizes: Option<Vec<String>>,
1342 ) -> Self {
1343 self.icons.get_or_insert_with(Vec::new).push(ToolIcon {
1344 src: src.into(),
1345 mime_type,
1346 sizes,
1347 theme: None,
1348 });
1349 self
1350 }
1351
1352 pub fn description(mut self, description: impl Into<String>) -> Self {
1354 self.description = Some(description.into());
1355 self
1356 }
1357
1358 pub fn read_only(mut self) -> Self {
1360 self.annotations
1361 .get_or_insert_with(ToolAnnotations::default)
1362 .read_only_hint = true;
1363 self
1364 }
1365
1366 pub fn non_destructive(mut self) -> Self {
1368 self.annotations
1369 .get_or_insert_with(ToolAnnotations::default)
1370 .destructive_hint = false;
1371 self
1372 }
1373
1374 pub fn destructive(mut self) -> Self {
1376 self.annotations
1377 .get_or_insert_with(ToolAnnotations::default)
1378 .destructive_hint = true;
1379 self
1380 }
1381
1382 pub fn idempotent(mut self) -> Self {
1384 self.annotations
1385 .get_or_insert_with(ToolAnnotations::default)
1386 .idempotent_hint = true;
1387 self
1388 }
1389
1390 pub fn read_only_safe(mut self) -> Self {
1396 let ann = self
1397 .annotations
1398 .get_or_insert_with(ToolAnnotations::default);
1399 ann.read_only_hint = true;
1400 ann.idempotent_hint = true;
1401 ann.destructive_hint = false;
1402 self
1403 }
1404
1405 pub fn annotations(mut self, annotations: ToolAnnotations) -> Self {
1407 self.annotations = Some(annotations);
1408 self
1409 }
1410
1411 pub fn task_support(mut self, mode: TaskSupportMode) -> Self {
1413 self.task_support = mode;
1414 self
1415 }
1416
1417 pub fn no_params_handler<F, Fut>(self, handler: F) -> ToolBuilderWithNoParamsHandler<F>
1435 where
1436 F: Fn() -> Fut + Send + Sync + 'static,
1437 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1438 {
1439 ToolBuilderWithNoParamsHandler {
1440 name: self.name,
1441 title: self.title,
1442 description: self.description,
1443 output_schema: self.output_schema,
1444 input_schema_override: self.input_schema_override,
1445 icons: self.icons,
1446 annotations: self.annotations,
1447 task_support: self.task_support,
1448 handler,
1449 }
1450 }
1451
1452 pub fn handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithHandler<I, F>
1495 where
1496 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1497 F: Fn(I) -> Fut + Send + Sync + 'static,
1498 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1499 {
1500 ToolBuilderWithHandler {
1501 name: self.name,
1502 title: self.title,
1503 description: self.description,
1504 output_schema: self.output_schema,
1505 input_schema_override: self.input_schema_override,
1506 icons: self.icons,
1507 annotations: self.annotations,
1508 task_support: self.task_support,
1509 task_preparer: None,
1510 handler,
1511 _phantom: std::marker::PhantomData,
1512 }
1513 }
1514
1515 #[cfg(feature = "stateless")]
1522 pub fn mrtr_handler<I, F, Fut>(self, handler: F) -> ToolBuilderWithMrtrHandler<I, F>
1523 where
1524 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1525 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
1526 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
1527 {
1528 ToolBuilderWithMrtrHandler {
1529 name: self.name,
1530 title: self.title,
1531 description: self.description,
1532 output_schema: self.output_schema,
1533 input_schema_override: self.input_schema_override,
1534 icons: self.icons,
1535 annotations: self.annotations,
1536 task_support: self.task_support,
1537 handler,
1538 _phantom: std::marker::PhantomData,
1539 }
1540 }
1541
1542 pub fn extractor_handler<S, F, T>(
1636 self,
1637 state: S,
1638 handler: F,
1639 ) -> crate::extract::ToolBuilderWithExtractor<S, F, T>
1640 where
1641 S: Clone + Send + Sync + 'static,
1642 F: crate::extract::ExtractorHandler<S, T> + Clone,
1643 T: Send + Sync + 'static,
1644 {
1645 let input_schema = ensure_object_schema(
1646 self.input_schema_override
1647 .unwrap_or_else(|| F::input_schema()),
1648 );
1649 crate::extract::ToolBuilderWithExtractor {
1650 name: self.name,
1651 title: self.title,
1652 description: self.description,
1653 output_schema: self.output_schema,
1654 icons: self.icons,
1655 annotations: self.annotations,
1656 task_support: self.task_support,
1657 state,
1658 handler,
1659 input_schema,
1660 _phantom: std::marker::PhantomData,
1661 }
1662 }
1663
1664 #[deprecated(
1698 since = "0.8.0",
1699 note = "Use `extractor_handler` instead -- it auto-detects JSON schema from `Json<T>` extractors without requiring a turbofish"
1700 )]
1701 #[allow(deprecated)]
1702 pub fn extractor_handler_typed<S, F, T, I>(
1703 self,
1704 state: S,
1705 handler: F,
1706 ) -> crate::extract::ToolBuilderWithTypedExtractor<S, F, T, I>
1707 where
1708 S: Clone + Send + Sync + 'static,
1709 F: crate::extract::TypedExtractorHandler<S, T, I> + Clone,
1710 T: Send + Sync + 'static,
1711 I: schemars::JsonSchema + Send + Sync + 'static,
1712 {
1713 crate::extract::ToolBuilderWithTypedExtractor {
1714 name: self.name,
1715 title: self.title,
1716 description: self.description,
1717 output_schema: self.output_schema,
1718 input_schema_override: self.input_schema_override,
1719 icons: self.icons,
1720 annotations: self.annotations,
1721 task_support: self.task_support,
1722 state,
1723 handler,
1724 _phantom: std::marker::PhantomData,
1725 }
1726 }
1727}
1728
1729struct NoParamsTypedHandler<F> {
1733 handler: F,
1734}
1735
1736impl<F, Fut> ToolHandler for NoParamsTypedHandler<F>
1737where
1738 F: Fn() -> Fut + Send + Sync + 'static,
1739 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1740{
1741 fn call(&self, _args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
1742 Box::pin(async move { (self.handler)().await })
1743 }
1744
1745 fn input_schema(&self) -> Value {
1746 serde_json::json!({ "type": "object" })
1747 }
1748}
1749
1750#[doc(hidden)]
1752pub struct ToolBuilderWithHandler<I, F> {
1753 name: String,
1754 title: Option<String>,
1755 description: Option<String>,
1756 output_schema: Option<Value>,
1757 input_schema_override: Option<Value>,
1758 icons: Option<Vec<ToolIcon>>,
1759 annotations: Option<ToolAnnotations>,
1760 task_support: TaskSupportMode,
1761 task_preparer: Option<Arc<dyn TaskPreparer>>,
1762 handler: F,
1763 _phantom: std::marker::PhantomData<I>,
1764}
1765
1766#[cfg(feature = "stateless")]
1768#[doc(hidden)]
1769pub struct ToolBuilderWithMrtrHandler<I, F> {
1770 name: String,
1771 title: Option<String>,
1772 description: Option<String>,
1773 output_schema: Option<Value>,
1774 input_schema_override: Option<Value>,
1775 icons: Option<Vec<ToolIcon>>,
1776 annotations: Option<ToolAnnotations>,
1777 task_support: TaskSupportMode,
1778 handler: F,
1779 _phantom: std::marker::PhantomData<I>,
1780}
1781
1782#[cfg(feature = "stateless")]
1784#[doc(hidden)]
1785pub struct ToolBuilderWithMrtrLayer<I, F, L> {
1786 name: String,
1787 title: Option<String>,
1788 description: Option<String>,
1789 output_schema: Option<Value>,
1790 input_schema_override: Option<Value>,
1791 icons: Option<Vec<ToolIcon>>,
1792 annotations: Option<ToolAnnotations>,
1793 task_support: TaskSupportMode,
1794 handler: F,
1795 layer: L,
1796 _phantom: std::marker::PhantomData<I>,
1797}
1798
1799#[doc(hidden)]
1803pub struct ToolBuilderWithNoParamsHandler<F> {
1804 name: String,
1805 title: Option<String>,
1806 description: Option<String>,
1807 output_schema: Option<Value>,
1808 input_schema_override: Option<Value>,
1809 icons: Option<Vec<ToolIcon>>,
1810 annotations: Option<ToolAnnotations>,
1811 task_support: TaskSupportMode,
1812 handler: F,
1813}
1814
1815impl<F, Fut> ToolBuilderWithNoParamsHandler<F>
1816where
1817 F: Fn() -> Fut + Send + Sync + 'static,
1818 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1819{
1820 pub fn build(self) -> Tool {
1822 Tool::from_handler(
1823 self.name,
1824 self.title,
1825 self.description,
1826 self.output_schema,
1827 self.icons,
1828 self.annotations,
1829 self.task_support,
1830 self.input_schema_override,
1831 NoParamsTypedHandler {
1832 handler: self.handler,
1833 },
1834 )
1835 }
1836
1837 pub fn layer<L>(self, layer: L) -> ToolBuilderWithNoParamsHandlerLayer<F, L> {
1841 ToolBuilderWithNoParamsHandlerLayer {
1842 name: self.name,
1843 title: self.title,
1844 description: self.description,
1845 output_schema: self.output_schema,
1846 input_schema_override: self.input_schema_override,
1847 icons: self.icons,
1848 annotations: self.annotations,
1849 task_support: self.task_support,
1850 handler: self.handler,
1851 layer,
1852 }
1853 }
1854
1855 pub fn guard<G>(self, guard: G) -> ToolBuilderWithNoParamsHandlerLayer<F, GuardLayer<G>>
1859 where
1860 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1861 {
1862 self.layer(GuardLayer::new(guard))
1863 }
1864}
1865
1866#[doc(hidden)]
1868pub struct ToolBuilderWithNoParamsHandlerLayer<F, L> {
1869 name: String,
1870 title: Option<String>,
1871 description: Option<String>,
1872 output_schema: Option<Value>,
1873 input_schema_override: Option<Value>,
1874 icons: Option<Vec<ToolIcon>>,
1875 annotations: Option<ToolAnnotations>,
1876 task_support: TaskSupportMode,
1877 handler: F,
1878 layer: L,
1879}
1880
1881#[allow(private_bounds)]
1882impl<F, Fut, L> ToolBuilderWithNoParamsHandlerLayer<F, L>
1883where
1884 F: Fn() -> Fut + Send + Sync + 'static,
1885 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1886 L: tower::Layer<ToolHandlerService<NoParamsTypedHandler<F>>> + Clone + Send + Sync + 'static,
1887 L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
1888 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
1889 <L::Service as Service<ToolRequest>>::Future: Send,
1890{
1891 pub fn build(self) -> Tool {
1893 let input_schema = ensure_object_schema(
1894 self.input_schema_override
1895 .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
1896 );
1897
1898 let handler_service = ToolHandlerService::new(NoParamsTypedHandler {
1899 handler: self.handler,
1900 });
1901 let layered = self.layer.layer(handler_service);
1902 let catch_error = ToolCatchError::new(layered);
1903 let service = BoxCloneService::new(catch_error);
1904
1905 Tool {
1906 name: self.name,
1907 title: self.title,
1908 description: self.description,
1909 output_schema: self.output_schema,
1910 icons: self.icons,
1911 annotations: self.annotations,
1912 meta: None,
1913 task_support: self.task_support,
1914 required_client_capabilities: None,
1915 task_preparer: None,
1916 service: Some(service),
1917 #[cfg(feature = "stateless")]
1918 mrtr_handler: None,
1919 input_schema,
1920 }
1921 }
1922
1923 pub fn layer<L2>(
1925 self,
1926 layer: L2,
1927 ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<L2, L>> {
1928 ToolBuilderWithNoParamsHandlerLayer {
1929 name: self.name,
1930 title: self.title,
1931 description: self.description,
1932 output_schema: self.output_schema,
1933 input_schema_override: self.input_schema_override,
1934 icons: self.icons,
1935 annotations: self.annotations,
1936 task_support: self.task_support,
1937 handler: self.handler,
1938 layer: tower::layer::util::Stack::new(layer, self.layer),
1939 }
1940 }
1941
1942 pub fn guard<G>(
1946 self,
1947 guard: G,
1948 ) -> ToolBuilderWithNoParamsHandlerLayer<F, tower::layer::util::Stack<GuardLayer<G>, L>>
1949 where
1950 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
1951 {
1952 self.layer(GuardLayer::new(guard))
1953 }
1954}
1955
1956impl<I, F, Fut> ToolBuilderWithHandler<I, F>
1957where
1958 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
1959 F: Fn(I) -> Fut + Send + Sync + 'static,
1960 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
1961{
1962 pub fn build(self) -> Tool {
1964 let mut tool = Tool::from_handler(
1965 self.name,
1966 self.title,
1967 self.description,
1968 self.output_schema,
1969 self.icons,
1970 self.annotations,
1971 self.task_support,
1972 self.input_schema_override,
1973 TypedHandler {
1974 handler: self.handler,
1975 _phantom: std::marker::PhantomData,
1976 },
1977 );
1978 tool.task_preparer = self.task_preparer;
1979 tool
1980 }
1981
1982 pub fn task_preparation<P, PrepareFuture>(mut self, prepare: P) -> Self
1984 where
1985 P: Fn(TaskContext, I) -> PrepareFuture + Send + Sync + 'static,
1986 PrepareFuture: Future<Output = Result<TaskPreparation>> + Send + 'static,
1987 {
1988 self.task_preparer = Some(Arc::new(TypedTaskPreparer {
1989 prepare,
1990 _phantom: std::marker::PhantomData,
1991 }));
1992 self
1993 }
1994
1995 pub fn layer<L>(self, layer: L) -> ToolBuilderWithLayer<I, F, L> {
2021 ToolBuilderWithLayer {
2022 name: self.name,
2023 title: self.title,
2024 description: self.description,
2025 output_schema: self.output_schema,
2026 input_schema_override: self.input_schema_override,
2027 icons: self.icons,
2028 annotations: self.annotations,
2029 task_support: self.task_support,
2030 task_preparer: self.task_preparer,
2031 handler: self.handler,
2032 layer,
2033 _phantom: std::marker::PhantomData,
2034 }
2035 }
2036
2037 pub fn guard<G>(self, guard: G) -> ToolBuilderWithLayer<I, F, GuardLayer<G>>
2044 where
2045 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2046 {
2047 self.layer(GuardLayer::new(guard))
2048 }
2049}
2050
2051#[cfg(feature = "stateless")]
2052impl<I, F, Fut> ToolBuilderWithMrtrHandler<I, F>
2053where
2054 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2055 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2056 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2057{
2058 pub fn build(self) -> Tool {
2060 Tool::from_mrtr_handler(
2061 self.name,
2062 self.title,
2063 self.description,
2064 self.output_schema,
2065 self.icons,
2066 self.annotations,
2067 self.task_support,
2068 self.input_schema_override,
2069 TypedMrtrHandler {
2070 handler: self.handler,
2071 _phantom: std::marker::PhantomData,
2072 },
2073 )
2074 }
2075
2076 pub fn layer<L>(self, layer: L) -> ToolBuilderWithMrtrLayer<I, F, L> {
2082 ToolBuilderWithMrtrLayer {
2083 name: self.name,
2084 title: self.title,
2085 description: self.description,
2086 output_schema: self.output_schema,
2087 input_schema_override: self.input_schema_override,
2088 icons: self.icons,
2089 annotations: self.annotations,
2090 task_support: self.task_support,
2091 handler: self.handler,
2092 layer,
2093 _phantom: std::marker::PhantomData,
2094 }
2095 }
2096
2097 pub fn guard<G>(self, guard: G) -> ToolBuilderWithMrtrLayer<I, F, GuardLayer<G>>
2099 where
2100 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2101 {
2102 self.layer(GuardLayer::new(guard))
2103 }
2104}
2105
2106#[cfg(feature = "stateless")]
2107#[allow(private_bounds)]
2108impl<I, F, Fut, L> ToolBuilderWithMrtrLayer<I, F, L>
2109where
2110 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2111 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2112 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2113 L: tower::Layer<MrtrToolHandlerService<TypedMrtrHandler<I, F>>> + Clone + Send + Sync + 'static,
2114 L::Service:
2115 Service<ToolRequest, Response = RequestOutcome<CallToolResult>> + Clone + Send + 'static,
2116 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send + 'static,
2117 <L::Service as Service<ToolRequest>>::Future: Send + 'static,
2118{
2119 pub fn build(self) -> Tool {
2121 let input_schema = self.input_schema_override.unwrap_or_else(|| {
2122 let schema = schemars::schema_for!(I);
2123 serde_json::to_value(schema).unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
2124 });
2125 let input_schema = ensure_object_schema(input_schema);
2126 let service = MrtrToolHandlerService::new(TypedMrtrHandler {
2127 handler: self.handler,
2128 _phantom: std::marker::PhantomData,
2129 });
2130 let service = self.layer.layer(service);
2131 let service = BoxCloneService::new(MrtrToolCatchError::new(service));
2132
2133 Tool {
2134 name: self.name,
2135 title: self.title,
2136 description: self.description,
2137 output_schema: self.output_schema,
2138 icons: self.icons,
2139 annotations: self.annotations,
2140 meta: None,
2141 task_support: self.task_support,
2142 required_client_capabilities: None,
2143 task_preparer: None,
2144 service: None,
2145 mrtr_handler: Some(Arc::new(ServiceMrtrToolHandler {
2146 service: Mutex::new(service),
2147 input_schema: input_schema.clone(),
2148 })),
2149 input_schema,
2150 }
2151 }
2152
2153 pub fn layer<L2>(
2155 self,
2156 layer: L2,
2157 ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<L2, L>> {
2158 ToolBuilderWithMrtrLayer {
2159 name: self.name,
2160 title: self.title,
2161 description: self.description,
2162 output_schema: self.output_schema,
2163 input_schema_override: self.input_schema_override,
2164 icons: self.icons,
2165 annotations: self.annotations,
2166 task_support: self.task_support,
2167 handler: self.handler,
2168 layer: tower::layer::util::Stack::new(layer, self.layer),
2169 _phantom: std::marker::PhantomData,
2170 }
2171 }
2172
2173 pub fn guard<G>(
2175 self,
2176 guard: G,
2177 ) -> ToolBuilderWithMrtrLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2178 where
2179 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2180 {
2181 self.layer(GuardLayer::new(guard))
2182 }
2183}
2184
2185#[doc(hidden)]
2189pub struct ToolBuilderWithLayer<I, F, L> {
2190 name: String,
2191 title: Option<String>,
2192 description: Option<String>,
2193 output_schema: Option<Value>,
2194 input_schema_override: Option<Value>,
2195 icons: Option<Vec<ToolIcon>>,
2196 annotations: Option<ToolAnnotations>,
2197 task_support: TaskSupportMode,
2198 task_preparer: Option<Arc<dyn TaskPreparer>>,
2199 handler: F,
2200 layer: L,
2201 _phantom: std::marker::PhantomData<I>,
2202}
2203
2204#[allow(private_bounds)]
2207impl<I, F, Fut, L> ToolBuilderWithLayer<I, F, L>
2208where
2209 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2210 F: Fn(I) -> Fut + Send + Sync + 'static,
2211 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2212 L: tower::Layer<ToolHandlerService<TypedHandler<I, F>>> + Clone + Send + Sync + 'static,
2213 L::Service: Service<ToolRequest, Response = CallToolResult> + Clone + Send + 'static,
2214 <L::Service as Service<ToolRequest>>::Error: fmt::Display + Send,
2215 <L::Service as Service<ToolRequest>>::Future: Send,
2216{
2217 pub fn build(self) -> Tool {
2219 let input_schema = self.input_schema_override.unwrap_or_else(|| {
2220 let input_schema = schemars::schema_for!(I);
2221 serde_json::to_value(input_schema)
2222 .unwrap_or_else(|_| serde_json::json!({ "type": "object" }))
2223 });
2224 let input_schema = ensure_object_schema(input_schema);
2225
2226 let handler_service = ToolHandlerService::new(TypedHandler {
2227 handler: self.handler,
2228 _phantom: std::marker::PhantomData,
2229 });
2230 let layered = self.layer.layer(handler_service);
2231 let catch_error = ToolCatchError::new(layered);
2232 let service = BoxCloneService::new(catch_error);
2233
2234 Tool {
2235 name: self.name,
2236 title: self.title,
2237 description: self.description,
2238 output_schema: self.output_schema,
2239 icons: self.icons,
2240 annotations: self.annotations,
2241 meta: None,
2242 task_support: self.task_support,
2243 required_client_capabilities: None,
2244 task_preparer: self.task_preparer,
2245 service: Some(service),
2246 #[cfg(feature = "stateless")]
2247 mrtr_handler: None,
2248 input_schema,
2249 }
2250 }
2251
2252 pub fn layer<L2>(
2257 self,
2258 layer: L2,
2259 ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<L2, L>> {
2260 ToolBuilderWithLayer {
2261 name: self.name,
2262 title: self.title,
2263 description: self.description,
2264 output_schema: self.output_schema,
2265 input_schema_override: self.input_schema_override,
2266 icons: self.icons,
2267 annotations: self.annotations,
2268 task_support: self.task_support,
2269 task_preparer: self.task_preparer,
2270 handler: self.handler,
2271 layer: tower::layer::util::Stack::new(layer, self.layer),
2272 _phantom: std::marker::PhantomData,
2273 }
2274 }
2275
2276 pub fn guard<G>(
2280 self,
2281 guard: G,
2282 ) -> ToolBuilderWithLayer<I, F, tower::layer::util::Stack<GuardLayer<G>, L>>
2283 where
2284 G: Fn(&ToolRequest) -> std::result::Result<(), String> + Clone + Send + Sync + 'static,
2285 {
2286 self.layer(GuardLayer::new(guard))
2287 }
2288}
2289
2290struct TypedHandler<I, F> {
2296 handler: F,
2297 _phantom: std::marker::PhantomData<I>,
2298}
2299
2300#[cfg(feature = "stateless")]
2301struct TypedMrtrHandler<I, F> {
2302 handler: F,
2303 _phantom: std::marker::PhantomData<I>,
2304}
2305
2306#[cfg(feature = "stateless")]
2307impl<I, F, Fut> MrtrToolHandler for TypedMrtrHandler<I, F>
2308where
2309 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2310 F: Fn(RequestContext, I) -> Fut + Send + Sync + 'static,
2311 Fut: Future<Output = Result<RequestOutcome<CallToolResult>>> + Send + 'static,
2312{
2313 fn call(
2314 &self,
2315 ctx: RequestContext,
2316 args: Value,
2317 ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
2318 Box::pin(async move {
2319 let input: I = serde_json::from_value(args)
2320 .map_err(|error| Error::invalid_params(format!("Invalid input: {error}")))?;
2321 (self.handler)(ctx, input).await
2322 })
2323 }
2324
2325 fn input_schema(&self) -> Value {
2326 let schema = schemars::schema_for!(I);
2327 ensure_object_schema(
2328 serde_json::to_value(schema)
2329 .unwrap_or_else(|_| serde_json::json!({ "type": "object" })),
2330 )
2331 }
2332}
2333
2334impl<I, F, Fut> ToolHandler for TypedHandler<I, F>
2335where
2336 I: JsonSchema + DeserializeOwned + Send + Sync + 'static,
2337 F: Fn(I) -> Fut + Send + Sync + 'static,
2338 Fut: Future<Output = Result<CallToolResult>> + Send + 'static,
2339{
2340 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2341 Box::pin(async move {
2342 let input: I = match serde_json::from_value(args) {
2343 Ok(input) => input,
2344 Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2345 };
2346 (self.handler)(input).await
2347 })
2348 }
2349
2350 fn input_schema(&self) -> Value {
2351 let schema = schemars::schema_for!(I);
2352 let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2353 serde_json::json!({
2354 "type": "object"
2355 })
2356 });
2357 ensure_object_schema(schema)
2358 }
2359}
2360
2361pub trait McpTool: Send + Sync + 'static {
2402 const NAME: &'static str;
2404 const DESCRIPTION: &'static str;
2406
2407 type Input: JsonSchema + DeserializeOwned + Send;
2409 type Output: Serialize + Send;
2411
2412 fn call(&self, input: Self::Input) -> impl Future<Output = Result<Self::Output>> + Send;
2414
2415 fn annotations(&self) -> Option<ToolAnnotations> {
2417 None
2418 }
2419
2420 fn into_tool(self) -> Tool
2428 where
2429 Self: Sized,
2430 {
2431 if let Err(e) = validate_tool_name(Self::NAME) {
2432 panic!("{e}");
2433 }
2434 let annotations = self.annotations();
2435 let tool = Arc::new(self);
2436 Tool::from_handler(
2437 Self::NAME.to_string(),
2438 None,
2439 Some(Self::DESCRIPTION.to_string()),
2440 None,
2441 None,
2442 annotations,
2443 TaskSupportMode::default(),
2444 None,
2445 McpToolHandler { tool },
2446 )
2447 }
2448}
2449
2450struct McpToolHandler<T: McpTool> {
2452 tool: Arc<T>,
2453}
2454
2455impl<T: McpTool> ToolHandler for McpToolHandler<T> {
2456 fn call(&self, args: Value) -> BoxFuture<'_, Result<CallToolResult>> {
2457 let tool = self.tool.clone();
2458 Box::pin(async move {
2459 let input: T::Input = match serde_json::from_value(args) {
2460 Ok(input) => input,
2461 Err(e) => return Ok(CallToolResult::error(format!("Invalid input: {e}"))),
2462 };
2463 let output = tool.call(input).await?;
2464 let value = serde_json::to_value(output).tool_context("Failed to serialize output")?;
2465 Ok(CallToolResult::json(value))
2466 })
2467 }
2468
2469 fn input_schema(&self) -> Value {
2470 let schema = schemars::schema_for!(T::Input);
2471 let schema = serde_json::to_value(schema).unwrap_or_else(|_| {
2472 serde_json::json!({
2473 "type": "object"
2474 })
2475 });
2476 ensure_object_schema(schema)
2477 }
2478}
2479
2480#[cfg(test)]
2481mod tests {
2482 use super::*;
2483 use crate::extract::{Context, Json, RawArgs, State};
2484 use crate::protocol::Content;
2485 use schemars::JsonSchema;
2486 use serde::Deserialize;
2487
2488 #[derive(Debug, Deserialize, JsonSchema)]
2489 struct GreetInput {
2490 name: String,
2491 }
2492
2493 #[tokio::test]
2494 async fn test_builder_tool() {
2495 let tool = ToolBuilder::new("greet")
2496 .description("Greet someone")
2497 .handler(|input: GreetInput| async move {
2498 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2499 })
2500 .build();
2501
2502 assert_eq!(tool.name, "greet");
2503 assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2504
2505 let result = tool.call(serde_json::json!({"name": "World"})).await;
2506
2507 assert!(!result.is_error);
2508 }
2509
2510 #[cfg(feature = "stateless")]
2511 #[tokio::test]
2512 async fn test_mrtr_builder_preserves_input_required_outcome() {
2513 let tool = ToolBuilder::new("continue")
2514 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2515 Ok(RequestOutcome::input_required(
2516 crate::protocol::InputRequiredResult::new().with_request_state("signed-state"),
2517 ))
2518 })
2519 .build();
2520
2521 let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2522 assert_eq!(
2523 outcome
2524 .as_input_required()
2525 .and_then(|result| result.request_state.as_deref()),
2526 Some("signed-state")
2527 );
2528 }
2529
2530 #[cfg(feature = "stateless")]
2531 #[tokio::test]
2532 async fn mrtr_builder_composes_guards_and_layers() {
2533 use std::sync::atomic::{AtomicUsize, Ordering};
2534 use std::time::Duration;
2535 use tower::timeout::TimeoutLayer;
2536
2537 let rounds = Arc::new(AtomicUsize::new(0));
2538 let observed = rounds.clone();
2539 let tool = ToolBuilder::new("guarded_continue")
2540 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2541 Ok(RequestOutcome::input_required(
2542 crate::protocol::InputRequiredResult::new().with_request_state("continue"),
2543 ))
2544 })
2545 .layer(TimeoutLayer::new(Duration::from_secs(1)))
2546 .guard(move |_request| {
2547 observed.fetch_add(1, Ordering::SeqCst);
2548 Ok(())
2549 })
2550 .build();
2551
2552 for _ in 0..2 {
2553 assert!(
2554 tool.call_outcome(serde_json::json!({}))
2555 .await
2556 .unwrap()
2557 .as_input_required()
2558 .is_some()
2559 );
2560 }
2561 assert_eq!(rounds.load(Ordering::SeqCst), 2);
2562 }
2563
2564 #[cfg(feature = "stateless")]
2565 #[tokio::test]
2566 async fn built_mrtr_tool_accepts_a_guard() {
2567 let tool = ToolBuilder::new("denied_continue")
2568 .mrtr_handler::<NoParams, _, _>(|_ctx, _input| async move {
2569 Ok(RequestOutcome::input_required(
2570 crate::protocol::InputRequiredResult::new().with_request_state("unreachable"),
2571 ))
2572 })
2573 .build()
2574 .with_guard(|_request| Err("MRTR access denied".to_string()));
2575
2576 let outcome = tool.call_outcome(serde_json::json!({})).await.unwrap();
2577 let result = outcome
2578 .as_complete()
2579 .expect("guard rejection is a complete tool error");
2580 assert!(result.is_error);
2581 assert_eq!(result.first_text(), Some("MRTR access denied"));
2582 }
2583
2584 #[tokio::test]
2585 async fn test_raw_handler() {
2586 let tool = ToolBuilder::new("echo")
2587 .description("Echo input")
2588 .extractor_handler((), |RawArgs(args): RawArgs| async move {
2589 Ok(CallToolResult::json(args))
2590 })
2591 .build();
2592
2593 let result = tool.call(serde_json::json!({"foo": "bar"})).await;
2594
2595 assert!(!result.is_error);
2596 }
2597
2598 #[test]
2599 fn test_invalid_tool_name_empty() {
2600 let err = ToolBuilder::try_new("").err().expect("should fail");
2601 assert!(err.to_string().contains("cannot be empty"));
2602 }
2603
2604 #[test]
2605 fn test_invalid_tool_name_too_long() {
2606 let long_name = "a".repeat(65);
2607 let err = ToolBuilder::try_new(long_name).err().expect("should fail");
2608 assert!(err.to_string().contains("exceeds maximum"));
2609 }
2610
2611 #[test]
2612 fn test_invalid_tool_name_bad_chars() {
2613 let err = ToolBuilder::try_new("my tool!").err().expect("should fail");
2614 assert!(err.to_string().contains("invalid character"));
2615 }
2616
2617 #[test]
2618 #[should_panic(expected = "cannot be empty")]
2619 fn test_new_panics_on_empty_name() {
2620 ToolBuilder::new("");
2621 }
2622
2623 #[test]
2624 #[should_panic(expected = "exceeds maximum")]
2625 fn test_new_panics_on_too_long_name() {
2626 ToolBuilder::new("a".repeat(65));
2627 }
2628
2629 #[test]
2630 #[should_panic(expected = "invalid character")]
2631 fn test_new_panics_on_invalid_chars() {
2632 ToolBuilder::new("my tool!");
2633 }
2634
2635 #[test]
2636 fn test_valid_tool_names() {
2637 let names = [
2639 "my_tool",
2640 "my-tool",
2641 "my.tool",
2642 "my/tool",
2643 "user-profile/update",
2644 "MyTool123",
2645 "a",
2646 &"a".repeat(64),
2647 ];
2648 for name in names {
2649 assert!(
2650 ToolBuilder::try_new(name).is_ok(),
2651 "Expected '{}' to be valid",
2652 name
2653 );
2654 }
2655 }
2656
2657 #[tokio::test]
2658 async fn test_context_aware_handler() {
2659 use crate::context::notification_channel;
2660 use crate::protocol::{ProgressToken, RequestId};
2661
2662 #[derive(Debug, Deserialize, JsonSchema)]
2663 struct ProcessInput {
2664 count: i32,
2665 }
2666
2667 let tool = ToolBuilder::new("process")
2668 .description("Process with context")
2669 .extractor_handler(
2670 (),
2671 |ctx: Context, Json(input): Json<ProcessInput>| async move {
2672 for i in 0..input.count {
2674 if ctx.is_cancelled() {
2675 return Ok(CallToolResult::error("Cancelled"));
2676 }
2677 ctx.report_progress(i as f64, Some(input.count as f64), None)
2678 .await;
2679 }
2680 Ok(CallToolResult::text(format!(
2681 "Processed {} items",
2682 input.count
2683 )))
2684 },
2685 )
2686 .build();
2687
2688 assert_eq!(tool.name, "process");
2689
2690 let (tx, mut rx) = notification_channel(10);
2692 let ctx = RequestContext::new(RequestId::Number(1))
2693 .with_progress_token(ProgressToken::Number(42))
2694 .with_notification_sender(tx);
2695
2696 let result = tool
2697 .call_with_context(ctx, serde_json::json!({"count": 3}))
2698 .await;
2699
2700 assert!(!result.is_error);
2701
2702 let mut progress_count = 0;
2704 while rx.try_recv().is_ok() {
2705 progress_count += 1;
2706 }
2707 assert_eq!(progress_count, 3);
2708 }
2709
2710 #[tokio::test]
2711 async fn test_context_aware_handler_cancellation() {
2712 use crate::protocol::RequestId;
2713 use std::sync::atomic::{AtomicI32, Ordering};
2714
2715 #[derive(Debug, Deserialize, JsonSchema)]
2716 struct LongRunningInput {
2717 iterations: i32,
2718 }
2719
2720 let iterations_completed = Arc::new(AtomicI32::new(0));
2721 let iterations_ref = iterations_completed.clone();
2722
2723 let tool = ToolBuilder::new("long_running")
2724 .description("Long running task")
2725 .extractor_handler(
2726 (),
2727 move |ctx: Context, Json(input): Json<LongRunningInput>| {
2728 let completed = iterations_ref.clone();
2729 async move {
2730 for i in 0..input.iterations {
2731 if ctx.is_cancelled() {
2732 return Ok(CallToolResult::error("Cancelled"));
2733 }
2734 completed.fetch_add(1, Ordering::SeqCst);
2735 tokio::task::yield_now().await;
2737 if i == 2 {
2739 ctx.cancellation_token().cancel();
2740 }
2741 }
2742 Ok(CallToolResult::text("Done"))
2743 }
2744 },
2745 )
2746 .build();
2747
2748 let ctx = RequestContext::new(RequestId::Number(1));
2749
2750 let result = tool
2751 .call_with_context(ctx, serde_json::json!({"iterations": 10}))
2752 .await;
2753
2754 assert!(result.is_error);
2757 assert_eq!(iterations_completed.load(Ordering::SeqCst), 3);
2758 }
2759
2760 #[tokio::test]
2761 async fn test_tool_builder_with_enhanced_fields() {
2762 let output_schema = serde_json::json!({
2763 "type": "object",
2764 "properties": {
2765 "greeting": {"type": "string"}
2766 }
2767 });
2768
2769 let tool = ToolBuilder::new("greet")
2770 .title("Greeting Tool")
2771 .description("Greet someone")
2772 .output_schema(output_schema.clone())
2773 .icon("https://example.com/icon.png")
2774 .icon_with_meta(
2775 "https://example.com/icon-large.png",
2776 Some("image/png".to_string()),
2777 Some(vec!["96x96".to_string()]),
2778 )
2779 .handler(|input: GreetInput| async move {
2780 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
2781 })
2782 .build();
2783
2784 assert_eq!(tool.name, "greet");
2785 assert_eq!(tool.title.as_deref(), Some("Greeting Tool"));
2786 assert_eq!(tool.description.as_deref(), Some("Greet someone"));
2787 assert_eq!(tool.output_schema, Some(output_schema));
2788 assert!(tool.icons.is_some());
2789 assert_eq!(tool.icons.as_ref().unwrap().len(), 2);
2790
2791 let def = tool.definition();
2793 assert_eq!(def.title.as_deref(), Some("Greeting Tool"));
2794 assert!(def.output_schema.is_some());
2795 assert!(def.icons.is_some());
2796 }
2797
2798 #[tokio::test]
2799 async fn test_handler_with_state() {
2800 let shared = Arc::new("shared-state".to_string());
2801
2802 let tool = ToolBuilder::new("stateful")
2803 .description("Uses shared state")
2804 .extractor_handler(
2805 shared,
2806 |State(state): State<Arc<String>>, Json(input): Json<GreetInput>| async move {
2807 Ok(CallToolResult::text(format!(
2808 "{}: Hello, {}!",
2809 state, input.name
2810 )))
2811 },
2812 )
2813 .build();
2814
2815 let result = tool.call(serde_json::json!({"name": "World"})).await;
2816 assert!(!result.is_error);
2817 }
2818
2819 #[tokio::test]
2820 async fn test_handler_with_state_and_context() {
2821 use crate::protocol::RequestId;
2822
2823 let shared = Arc::new(42_i32);
2824
2825 let tool =
2826 ToolBuilder::new("stateful_ctx")
2827 .description("Uses state and context")
2828 .extractor_handler(
2829 shared,
2830 |State(state): State<Arc<i32>>,
2831 _ctx: Context,
2832 Json(input): Json<GreetInput>| async move {
2833 Ok(CallToolResult::text(format!(
2834 "{}: Hello, {}!",
2835 state, input.name
2836 )))
2837 },
2838 )
2839 .build();
2840
2841 let ctx = RequestContext::new(RequestId::Number(1));
2842 let result = tool
2843 .call_with_context(ctx, serde_json::json!({"name": "World"}))
2844 .await;
2845 assert!(!result.is_error);
2846 }
2847
2848 #[tokio::test]
2849 async fn test_handler_no_params() {
2850 let tool = ToolBuilder::new("no_params")
2851 .description("Takes no parameters")
2852 .extractor_handler((), |Json(_): Json<NoParams>| async {
2853 Ok(CallToolResult::text("no params result"))
2854 })
2855 .build();
2856
2857 assert_eq!(tool.name, "no_params");
2858
2859 let result = tool.call(serde_json::json!({})).await;
2861 assert!(!result.is_error);
2862
2863 let result = tool.call(serde_json::json!({"unexpected": "value"})).await;
2865 assert!(!result.is_error);
2866
2867 let schema = tool.definition().input_schema;
2869 assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2870 }
2871
2872 #[tokio::test]
2873 async fn test_handler_with_state_no_params() {
2874 let shared = Arc::new("shared_value".to_string());
2875
2876 let tool = ToolBuilder::new("with_state_no_params")
2877 .description("Takes no parameters but has state")
2878 .extractor_handler(
2879 shared,
2880 |State(state): State<Arc<String>>, Json(_): Json<NoParams>| async move {
2881 Ok(CallToolResult::text(format!("state: {}", state)))
2882 },
2883 )
2884 .build();
2885
2886 assert_eq!(tool.name, "with_state_no_params");
2887
2888 let result = tool.call(serde_json::json!({})).await;
2890 assert!(!result.is_error);
2891 assert_eq!(result.first_text().unwrap(), "state: shared_value");
2892
2893 let schema = tool.definition().input_schema;
2895 assert_eq!(schema.get("type").unwrap().as_str().unwrap(), "object");
2896 }
2897
2898 #[tokio::test]
2899 async fn test_handler_no_params_with_context() {
2900 let tool = ToolBuilder::new("no_params_with_context")
2901 .description("Takes no parameters but has context")
2902 .extractor_handler((), |_ctx: Context, Json(_): Json<NoParams>| async move {
2903 Ok(CallToolResult::text("context available"))
2904 })
2905 .build();
2906
2907 assert_eq!(tool.name, "no_params_with_context");
2908
2909 let result = tool.call(serde_json::json!({})).await;
2910 assert!(!result.is_error);
2911 assert_eq!(result.first_text().unwrap(), "context available");
2912 }
2913
2914 #[tokio::test]
2915 async fn test_handler_with_state_and_context_no_params() {
2916 let shared = Arc::new("shared".to_string());
2917
2918 let tool = ToolBuilder::new("state_context_no_params")
2919 .description("Has state and context, no params")
2920 .extractor_handler(
2921 shared,
2922 |State(state): State<Arc<String>>,
2923 _ctx: Context,
2924 Json(_): Json<NoParams>| async move {
2925 Ok(CallToolResult::text(format!("state: {}", state)))
2926 },
2927 )
2928 .build();
2929
2930 assert_eq!(tool.name, "state_context_no_params");
2931
2932 let result = tool.call(serde_json::json!({})).await;
2933 assert!(!result.is_error);
2934 assert_eq!(result.first_text().unwrap(), "state: shared");
2935 }
2936
2937 #[tokio::test]
2938 async fn test_raw_handler_with_state() {
2939 let prefix = Arc::new("prefix:".to_string());
2940
2941 let tool = ToolBuilder::new("raw_with_state")
2942 .description("Raw handler with state")
2943 .extractor_handler(
2944 prefix,
2945 |State(state): State<Arc<String>>, RawArgs(args): RawArgs| async move {
2946 Ok(CallToolResult::text(format!("{} {}", state, args)))
2947 },
2948 )
2949 .build();
2950
2951 assert_eq!(tool.name, "raw_with_state");
2952
2953 let result = tool.call(serde_json::json!({"key": "value"})).await;
2954 assert!(!result.is_error);
2955 assert!(result.first_text().unwrap().starts_with("prefix:"));
2956 }
2957
2958 #[tokio::test]
2959 async fn test_raw_handler_with_state_and_context() {
2960 let prefix = Arc::new("prefix:".to_string());
2961
2962 let tool = ToolBuilder::new("raw_state_context")
2963 .description("Raw handler with state and context")
2964 .extractor_handler(
2965 prefix,
2966 |State(state): State<Arc<String>>,
2967 _ctx: Context,
2968 RawArgs(args): RawArgs| async move {
2969 Ok(CallToolResult::text(format!("{} {}", state, args)))
2970 },
2971 )
2972 .build();
2973
2974 assert_eq!(tool.name, "raw_state_context");
2975
2976 let result = tool.call(serde_json::json!({"key": "value"})).await;
2977 assert!(!result.is_error);
2978 assert!(result.first_text().unwrap().starts_with("prefix:"));
2979 }
2980
2981 #[tokio::test]
2982 async fn test_tool_with_timeout_layer() {
2983 use std::time::Duration;
2984 use tower::timeout::TimeoutLayer;
2985
2986 #[derive(Debug, Deserialize, JsonSchema)]
2987 struct SlowInput {
2988 delay_ms: u64,
2989 }
2990
2991 let tool = ToolBuilder::new("slow_tool")
2993 .description("A slow tool")
2994 .handler(|input: SlowInput| async move {
2995 tokio::time::sleep(Duration::from_millis(input.delay_ms)).await;
2996 Ok(CallToolResult::text("completed"))
2997 })
2998 .layer(TimeoutLayer::new(Duration::from_millis(50)))
2999 .build();
3000
3001 let result = tool.call(serde_json::json!({"delay_ms": 10})).await;
3003 assert!(!result.is_error);
3004 assert_eq!(result.first_text().unwrap(), "completed");
3005
3006 let result = tool.call(serde_json::json!({"delay_ms": 200})).await;
3008 assert!(result.is_error);
3009 let msg = result.first_text().unwrap().to_lowercase();
3011 assert!(
3012 msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
3013 "Expected timeout error, got: {}",
3014 msg
3015 );
3016 }
3017
3018 #[tokio::test]
3019 async fn test_tool_with_concurrency_limit_layer() {
3020 use std::sync::atomic::{AtomicU32, Ordering};
3021 use std::time::Duration;
3022 use tower::limit::ConcurrencyLimitLayer;
3023
3024 #[derive(Debug, Deserialize, JsonSchema)]
3025 struct WorkInput {
3026 id: u32,
3027 }
3028
3029 let max_concurrent = Arc::new(AtomicU32::new(0));
3030 let current_concurrent = Arc::new(AtomicU32::new(0));
3031 let max_ref = max_concurrent.clone();
3032 let current_ref = current_concurrent.clone();
3033
3034 let tool = ToolBuilder::new("concurrent_tool")
3036 .description("A concurrent tool")
3037 .handler(move |input: WorkInput| {
3038 let max = max_ref.clone();
3039 let current = current_ref.clone();
3040 async move {
3041 let prev = current.fetch_add(1, Ordering::SeqCst);
3043 max.fetch_max(prev + 1, Ordering::SeqCst);
3044
3045 tokio::time::sleep(Duration::from_millis(50)).await;
3047
3048 current.fetch_sub(1, Ordering::SeqCst);
3049 Ok(CallToolResult::text(format!("completed {}", input.id)))
3050 }
3051 })
3052 .layer(ConcurrencyLimitLayer::new(2))
3053 .build();
3054
3055 let handles: Vec<_> = (0..4)
3057 .map(|i| {
3058 let t = tool.call(serde_json::json!({"id": i}));
3059 tokio::spawn(t)
3060 })
3061 .collect();
3062
3063 for handle in handles {
3064 let result = handle.await.unwrap();
3065 assert!(!result.is_error);
3066 }
3067
3068 assert!(max_concurrent.load(Ordering::SeqCst) <= 2);
3070 }
3071
3072 #[tokio::test]
3073 async fn test_tool_with_multiple_layers() {
3074 use std::time::Duration;
3075 use tower::limit::ConcurrencyLimitLayer;
3076 use tower::timeout::TimeoutLayer;
3077
3078 #[derive(Debug, Deserialize, JsonSchema)]
3079 struct Input {
3080 value: String,
3081 }
3082
3083 let tool = ToolBuilder::new("multi_layer_tool")
3085 .description("Tool with multiple layers")
3086 .handler(|input: Input| async move {
3087 Ok(CallToolResult::text(format!("processed: {}", input.value)))
3088 })
3089 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3090 .layer(ConcurrencyLimitLayer::new(10))
3091 .build();
3092
3093 let result = tool.call(serde_json::json!({"value": "test"})).await;
3094 assert!(!result.is_error);
3095 assert_eq!(result.first_text().unwrap(), "processed: test");
3096 }
3097
3098 #[test]
3099 fn test_tool_catch_error_clone() {
3100 let tool = ToolBuilder::new("test")
3103 .description("test")
3104 .extractor_handler((), |RawArgs(_args): RawArgs| async {
3105 Ok(CallToolResult::text("ok"))
3106 })
3107 .build();
3108 let _clone = tool.call(serde_json::json!({}));
3110 }
3111
3112 #[test]
3113 fn test_tool_catch_error_debug() {
3114 #[derive(Debug, Clone)]
3118 struct DebugService;
3119
3120 impl Service<ToolRequest> for DebugService {
3121 type Response = CallToolResult;
3122 type Error = crate::error::Error;
3123 type Future = Pin<
3124 Box<
3125 dyn Future<Output = std::result::Result<CallToolResult, crate::error::Error>>
3126 + Send,
3127 >,
3128 >;
3129
3130 fn poll_ready(
3131 &mut self,
3132 _cx: &mut std::task::Context<'_>,
3133 ) -> Poll<std::result::Result<(), Self::Error>> {
3134 Poll::Ready(Ok(()))
3135 }
3136
3137 fn call(&mut self, _req: ToolRequest) -> Self::Future {
3138 Box::pin(async { Ok(CallToolResult::text("ok")) })
3139 }
3140 }
3141
3142 let catch_error = ToolCatchError::new(DebugService);
3143 let debug = format!("{:?}", catch_error);
3144 assert!(debug.contains("ToolCatchError"));
3145 }
3146
3147 #[test]
3148 fn test_tool_request_new() {
3149 use crate::protocol::RequestId;
3150
3151 let ctx = RequestContext::new(RequestId::Number(42));
3152 let args = serde_json::json!({"key": "value"});
3153 let req = ToolRequest::new(ctx.clone(), args.clone());
3154
3155 assert_eq!(req.args, args);
3156 }
3157
3158 #[test]
3159 fn test_no_params_schema() {
3160 let schema = schemars::schema_for!(NoParams);
3162 let schema_value = serde_json::to_value(&schema).unwrap();
3163 assert_eq!(
3164 schema_value.get("type").and_then(|v| v.as_str()),
3165 Some("object"),
3166 "NoParams should generate type: object schema"
3167 );
3168 }
3169
3170 #[test]
3171 fn test_no_params_deserialize() {
3172 let from_empty_object: NoParams = serde_json::from_str("{}").unwrap();
3174 assert_eq!(from_empty_object, NoParams);
3175
3176 let from_null: NoParams = serde_json::from_str("null").unwrap();
3177 assert_eq!(from_null, NoParams);
3178
3179 let from_object_with_fields: NoParams =
3181 serde_json::from_str(r#"{"unexpected": "value"}"#).unwrap();
3182 assert_eq!(from_object_with_fields, NoParams);
3183 }
3184
3185 #[tokio::test]
3186 async fn test_no_params_type_in_handler() {
3187 let tool = ToolBuilder::new("status")
3189 .description("Get status")
3190 .handler(|_input: NoParams| async move { Ok(CallToolResult::text("OK")) })
3191 .build();
3192
3193 let schema = tool.definition().input_schema;
3195 assert_eq!(
3196 schema.get("type").and_then(|v| v.as_str()),
3197 Some("object"),
3198 "NoParams handler should produce type: object schema"
3199 );
3200
3201 let result = tool.call(serde_json::json!({})).await;
3203 assert!(!result.is_error);
3204 }
3205
3206 #[tokio::test]
3207 async fn test_serde_json_value_handler_has_type_object() {
3208 let tool = ToolBuilder::new("any_input")
3211 .description("Accepts any input")
3212 .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3213 .build();
3214
3215 let schema = tool.definition().input_schema;
3216 assert_eq!(
3217 schema.get("type").and_then(|v| v.as_str()),
3218 Some("object"),
3219 "serde_json::Value handler should produce schema with type: object"
3220 );
3221 }
3222
3223 #[tokio::test]
3224 async fn test_tool_with_name_prefix() {
3225 #[derive(Debug, Deserialize, JsonSchema)]
3226 struct Input {
3227 value: String,
3228 }
3229
3230 let tool = ToolBuilder::new("query")
3231 .description("Query something")
3232 .title("Query Tool")
3233 .handler(|input: Input| async move { Ok(CallToolResult::text(&input.value)) })
3234 .build();
3235
3236 let prefixed = tool.with_name_prefix("db");
3238
3239 assert_eq!(prefixed.name, "db.query");
3241
3242 assert_eq!(prefixed.description.as_deref(), Some("Query something"));
3244 assert_eq!(prefixed.title.as_deref(), Some("Query Tool"));
3245
3246 let result = prefixed
3248 .call(serde_json::json!({"value": "test input"}))
3249 .await;
3250 assert!(!result.is_error);
3251 match &result.content[0] {
3252 Content::Text { text, .. } => assert_eq!(text, "test input"),
3253 _ => panic!("Expected text content"),
3254 }
3255 }
3256
3257 #[tokio::test]
3258 async fn test_tool_with_name_prefix_multiple_levels() {
3259 let tool = ToolBuilder::new("action")
3260 .description("Do something")
3261 .handler(|_: NoParams| async move { Ok(CallToolResult::text("done")) })
3262 .build();
3263
3264 let prefixed = tool.with_name_prefix("level1");
3266 assert_eq!(prefixed.name, "level1.action");
3267
3268 let double_prefixed = prefixed.with_name_prefix("level0");
3269 assert_eq!(double_prefixed.name, "level0.level1.action");
3270 }
3271
3272 #[tokio::test]
3277 async fn test_no_params_handler_basic() {
3278 let tool = ToolBuilder::new("get_status")
3279 .description("Get current status")
3280 .no_params_handler(|| async { Ok(CallToolResult::text("OK")) })
3281 .build();
3282
3283 assert_eq!(tool.name, "get_status");
3284 assert_eq!(tool.description.as_deref(), Some("Get current status"));
3285
3286 let result = tool.call(serde_json::json!({})).await;
3288 assert!(!result.is_error);
3289 assert_eq!(result.first_text().unwrap(), "OK");
3290
3291 let result = tool.call(serde_json::json!(null)).await;
3293 assert!(!result.is_error);
3294
3295 let schema = tool.definition().input_schema;
3297 assert_eq!(schema.get("type").and_then(|v| v.as_str()), Some("object"));
3298 }
3299
3300 #[tokio::test]
3301 async fn test_no_params_handler_with_captured_state() {
3302 let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
3303 let counter_ref = counter.clone();
3304
3305 let tool = ToolBuilder::new("increment")
3306 .description("Increment counter")
3307 .no_params_handler(move || {
3308 let c = counter_ref.clone();
3309 async move {
3310 let prev = c.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
3311 Ok(CallToolResult::text(format!("Incremented from {}", prev)))
3312 }
3313 })
3314 .build();
3315
3316 let _ = tool.call(serde_json::json!({})).await;
3318 let _ = tool.call(serde_json::json!({})).await;
3319 let result = tool.call(serde_json::json!({})).await;
3320
3321 assert!(!result.is_error);
3322 assert_eq!(result.first_text().unwrap(), "Incremented from 2");
3323 assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 3);
3324 }
3325
3326 #[tokio::test]
3327 async fn test_no_params_handler_with_layer() {
3328 use std::time::Duration;
3329 use tower::timeout::TimeoutLayer;
3330
3331 let tool = ToolBuilder::new("slow_status")
3332 .description("Slow status check")
3333 .no_params_handler(|| async {
3334 tokio::time::sleep(Duration::from_millis(10)).await;
3335 Ok(CallToolResult::text("done"))
3336 })
3337 .layer(TimeoutLayer::new(Duration::from_secs(1)))
3338 .build();
3339
3340 let result = tool.call(serde_json::json!({})).await;
3341 assert!(!result.is_error);
3342 assert_eq!(result.first_text().unwrap(), "done");
3343 }
3344
3345 #[tokio::test]
3346 async fn test_no_params_handler_timeout() {
3347 use std::time::Duration;
3348 use tower::timeout::TimeoutLayer;
3349
3350 let tool = ToolBuilder::new("very_slow_status")
3351 .description("Very slow status check")
3352 .no_params_handler(|| async {
3353 tokio::time::sleep(Duration::from_millis(200)).await;
3354 Ok(CallToolResult::text("done"))
3355 })
3356 .layer(TimeoutLayer::new(Duration::from_millis(50)))
3357 .build();
3358
3359 let result = tool.call(serde_json::json!({})).await;
3360 assert!(result.is_error);
3361 let msg = result.first_text().unwrap().to_lowercase();
3362 assert!(
3363 msg.contains("timed out") || msg.contains("timeout") || msg.contains("elapsed"),
3364 "Expected timeout error, got: {}",
3365 msg
3366 );
3367 }
3368
3369 #[tokio::test]
3370 async fn test_no_params_handler_with_multiple_layers() {
3371 use std::time::Duration;
3372 use tower::limit::ConcurrencyLimitLayer;
3373 use tower::timeout::TimeoutLayer;
3374
3375 let tool = ToolBuilder::new("multi_layer_status")
3376 .description("Status with multiple layers")
3377 .no_params_handler(|| async { Ok(CallToolResult::text("status ok")) })
3378 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3379 .layer(ConcurrencyLimitLayer::new(10))
3380 .build();
3381
3382 let result = tool.call(serde_json::json!({})).await;
3383 assert!(!result.is_error);
3384 assert_eq!(result.first_text().unwrap(), "status ok");
3385 }
3386
3387 #[tokio::test]
3392 async fn test_guard_allows_request() {
3393 #[derive(Debug, Deserialize, JsonSchema)]
3394 #[allow(dead_code)]
3395 struct DeleteInput {
3396 id: String,
3397 confirm: bool,
3398 }
3399
3400 let tool = ToolBuilder::new("delete")
3401 .description("Delete a record")
3402 .handler(|input: DeleteInput| async move {
3403 Ok(CallToolResult::text(format!("deleted {}", input.id)))
3404 })
3405 .guard(|req: &ToolRequest| {
3406 let confirm = req
3407 .args
3408 .get("confirm")
3409 .and_then(|v| v.as_bool())
3410 .unwrap_or(false);
3411 if !confirm {
3412 return Err("Must set confirm=true to delete".to_string());
3413 }
3414 Ok(())
3415 })
3416 .build();
3417
3418 let result = tool
3419 .call(serde_json::json!({"id": "abc", "confirm": true}))
3420 .await;
3421 assert!(!result.is_error);
3422 assert_eq!(result.first_text().unwrap(), "deleted abc");
3423 }
3424
3425 #[tokio::test]
3426 async fn test_guard_rejects_request() {
3427 #[derive(Debug, Deserialize, JsonSchema)]
3428 #[allow(dead_code)]
3429 struct DeleteInput2 {
3430 id: String,
3431 confirm: bool,
3432 }
3433
3434 let tool = ToolBuilder::new("delete2")
3435 .description("Delete a record")
3436 .handler(|input: DeleteInput2| async move {
3437 Ok(CallToolResult::text(format!("deleted {}", input.id)))
3438 })
3439 .guard(|req: &ToolRequest| {
3440 let confirm = req
3441 .args
3442 .get("confirm")
3443 .and_then(|v| v.as_bool())
3444 .unwrap_or(false);
3445 if !confirm {
3446 return Err("Must set confirm=true to delete".to_string());
3447 }
3448 Ok(())
3449 })
3450 .build();
3451
3452 let result = tool
3453 .call(serde_json::json!({"id": "abc", "confirm": false}))
3454 .await;
3455 assert!(result.is_error);
3456 assert!(
3457 result
3458 .first_text()
3459 .unwrap()
3460 .contains("Must set confirm=true")
3461 );
3462 }
3463
3464 #[tokio::test]
3465 async fn test_guard_with_layer() {
3466 use std::time::Duration;
3467 use tower::timeout::TimeoutLayer;
3468
3469 let tool = ToolBuilder::new("guarded_timeout")
3470 .description("Guarded with timeout")
3471 .handler(|input: GreetInput| async move {
3472 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3473 })
3474 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3475 .guard(|_req: &ToolRequest| Ok(()))
3476 .build();
3477
3478 let result = tool.call(serde_json::json!({"name": "World"})).await;
3479 assert!(!result.is_error);
3480 assert_eq!(result.first_text().unwrap(), "Hello, World!");
3481 }
3482
3483 #[tokio::test]
3484 async fn test_guard_on_no_params_handler() {
3485 let allowed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
3486 let allowed_clone = allowed.clone();
3487
3488 let tool = ToolBuilder::new("status")
3489 .description("Get status")
3490 .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3491 .guard(move |_req: &ToolRequest| {
3492 if allowed_clone.load(std::sync::atomic::Ordering::Relaxed) {
3493 Ok(())
3494 } else {
3495 Err("Access denied".to_string())
3496 }
3497 })
3498 .build();
3499
3500 let result = tool.call(serde_json::json!({})).await;
3502 assert!(!result.is_error);
3503 assert_eq!(result.first_text().unwrap(), "ok");
3504
3505 allowed.store(false, std::sync::atomic::Ordering::Relaxed);
3507 let result = tool.call(serde_json::json!({})).await;
3508 assert!(result.is_error);
3509 assert!(result.first_text().unwrap().contains("Access denied"));
3510 }
3511
3512 #[tokio::test]
3513 async fn test_guard_on_no_params_handler_with_layer() {
3514 use std::time::Duration;
3515 use tower::timeout::TimeoutLayer;
3516
3517 let tool = ToolBuilder::new("status_layered")
3518 .description("Get status with layers")
3519 .no_params_handler(|| async { Ok(CallToolResult::text("ok")) })
3520 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3521 .guard(|_req: &ToolRequest| Ok(()))
3522 .build();
3523
3524 let result = tool.call(serde_json::json!({})).await;
3525 assert!(!result.is_error);
3526 assert_eq!(result.first_text().unwrap(), "ok");
3527 }
3528
3529 #[tokio::test]
3530 async fn test_guard_on_extractor_handler() {
3531 use std::sync::Arc;
3532
3533 #[derive(Clone)]
3534 struct AppState {
3535 prefix: String,
3536 }
3537
3538 #[derive(Debug, Deserialize, JsonSchema)]
3539 struct QueryInput {
3540 query: String,
3541 }
3542
3543 let state = Arc::new(AppState {
3544 prefix: "db".to_string(),
3545 });
3546
3547 let tool = ToolBuilder::new("search")
3548 .description("Search")
3549 .extractor_handler(
3550 state,
3551 |State(app): State<Arc<AppState>>, Json(input): Json<QueryInput>| async move {
3552 Ok(CallToolResult::text(format!(
3553 "{}: {}",
3554 app.prefix, input.query
3555 )))
3556 },
3557 )
3558 .guard(|req: &ToolRequest| {
3559 let query = req.args.get("query").and_then(|v| v.as_str()).unwrap_or("");
3560 if query.is_empty() {
3561 return Err("Query cannot be empty".to_string());
3562 }
3563 Ok(())
3564 })
3565 .build();
3566
3567 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3569 assert!(!result.is_error);
3570 assert_eq!(result.first_text().unwrap(), "db: hello");
3571
3572 let result = tool.call(serde_json::json!({"query": ""})).await;
3574 assert!(result.is_error);
3575 assert!(
3576 result
3577 .first_text()
3578 .unwrap()
3579 .contains("Query cannot be empty")
3580 );
3581 }
3582
3583 #[tokio::test]
3584 async fn test_guard_on_extractor_handler_with_layer() {
3585 use std::sync::Arc;
3586 use std::time::Duration;
3587 use tower::timeout::TimeoutLayer;
3588
3589 #[derive(Clone)]
3590 struct AppState2 {
3591 prefix: String,
3592 }
3593
3594 #[derive(Debug, Deserialize, JsonSchema)]
3595 struct QueryInput2 {
3596 query: String,
3597 }
3598
3599 let state = Arc::new(AppState2 {
3600 prefix: "db".to_string(),
3601 });
3602
3603 let tool = ToolBuilder::new("search2")
3604 .description("Search with layer and guard")
3605 .extractor_handler(
3606 state,
3607 |State(app): State<Arc<AppState2>>, Json(input): Json<QueryInput2>| async move {
3608 Ok(CallToolResult::text(format!(
3609 "{}: {}",
3610 app.prefix, input.query
3611 )))
3612 },
3613 )
3614 .layer(TimeoutLayer::new(Duration::from_secs(5)))
3615 .guard(|_req: &ToolRequest| Ok(()))
3616 .build();
3617
3618 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3619 assert!(!result.is_error);
3620 assert_eq!(result.first_text().unwrap(), "db: hello");
3621 }
3622
3623 #[tokio::test]
3624 async fn test_tool_with_guard_post_build() {
3625 let tool = ToolBuilder::new("admin_action")
3626 .description("Admin action")
3627 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3628 .build();
3629
3630 let guarded = tool.with_guard(|req: &ToolRequest| {
3632 let name = req.args.get("name").and_then(|v| v.as_str()).unwrap_or("");
3633 if name == "admin" {
3634 Ok(())
3635 } else {
3636 Err("Only admin allowed".to_string())
3637 }
3638 });
3639
3640 let result = guarded.call(serde_json::json!({"name": "admin"})).await;
3642 assert!(!result.is_error);
3643
3644 let result = guarded.call(serde_json::json!({"name": "user"})).await;
3646 assert!(result.is_error);
3647 assert!(result.first_text().unwrap().contains("Only admin allowed"));
3648 }
3649
3650 #[tokio::test]
3651 async fn test_with_guard_preserves_tool_metadata() {
3652 let tool = ToolBuilder::new("my_tool")
3653 .description("A tool")
3654 .title("My Tool")
3655 .read_only()
3656 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("done")) })
3657 .build();
3658
3659 let guarded = tool.with_guard(|_req: &ToolRequest| Ok(()));
3660
3661 assert_eq!(guarded.name, "my_tool");
3662 assert_eq!(guarded.description.as_deref(), Some("A tool"));
3663 assert_eq!(guarded.title.as_deref(), Some("My Tool"));
3664 assert!(guarded.annotations.is_some());
3665 }
3666
3667 #[tokio::test]
3668 async fn test_guard_group_pattern() {
3669 let require_auth = |req: &ToolRequest| {
3671 let token = req
3672 .args
3673 .get("_token")
3674 .and_then(|v| v.as_str())
3675 .unwrap_or("");
3676 if token == "valid" {
3677 Ok(())
3678 } else {
3679 Err("Authentication required".to_string())
3680 }
3681 };
3682
3683 let tool1 = ToolBuilder::new("action1")
3684 .description("Action 1")
3685 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action1")) })
3686 .build();
3687 let tool2 = ToolBuilder::new("action2")
3688 .description("Action 2")
3689 .handler(|_input: GreetInput| async move { Ok(CallToolResult::text("action2")) })
3690 .build();
3691
3692 let guarded1 = tool1.with_guard(require_auth);
3694 let guarded2 = tool2.with_guard(require_auth);
3695
3696 let r1 = guarded1
3698 .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3699 .await;
3700 let r2 = guarded2
3701 .call(serde_json::json!({"name": "test", "_token": "invalid"}))
3702 .await;
3703 assert!(r1.is_error);
3704 assert!(r2.is_error);
3705
3706 let r1 = guarded1
3708 .call(serde_json::json!({"name": "test", "_token": "valid"}))
3709 .await;
3710 let r2 = guarded2
3711 .call(serde_json::json!({"name": "test", "_token": "valid"}))
3712 .await;
3713 assert!(!r1.is_error);
3714 assert!(!r2.is_error);
3715 }
3716
3717 #[tokio::test]
3718 async fn test_input_validation_returns_tool_error() {
3719 #[derive(Debug, Deserialize, JsonSchema)]
3722 struct StrictInput {
3723 name: String,
3724 count: u32,
3725 }
3726
3727 let tool = ToolBuilder::new("strict_tool")
3728 .description("requires specific input")
3729 .handler(|input: StrictInput| async move {
3730 Ok(CallToolResult::text(format!(
3731 "{}: {}",
3732 input.name, input.count
3733 )))
3734 })
3735 .build();
3736
3737 let result = tool
3739 .call(serde_json::json!({"name": "test", "count": 5}))
3740 .await;
3741 assert!(!result.is_error);
3742
3743 let result = tool.call(serde_json::json!({"name": "test"})).await;
3745 assert!(result.is_error);
3746 let text = result.first_text().unwrap();
3747 assert!(text.contains("Invalid input"), "got: {text}");
3748
3749 let result = tool
3751 .call(serde_json::json!({"name": "test", "count": "not_a_number"}))
3752 .await;
3753 assert!(result.is_error);
3754 let text = result.first_text().unwrap();
3755 assert!(text.contains("Invalid input"), "got: {text}");
3756 }
3757
3758 #[tokio::test]
3759 async fn test_input_schema_override_with_raw_args() {
3760 let custom = serde_json::json!({
3764 "type": "object",
3765 "properties": {
3766 "query": { "type": "string", "minLength": 1 }
3767 },
3768 "required": ["query"]
3769 });
3770
3771 let tool = ToolBuilder::new("query")
3772 .description("Query with a custom schema")
3773 .input_schema(custom.clone())
3774 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3775 Ok(CallToolResult::json(args))
3776 })
3777 .build();
3778
3779 let schema = tool.definition().input_schema;
3780 assert_eq!(schema, custom);
3781
3782 let result = tool.call(serde_json::json!({"query": "hello"})).await;
3784 assert!(!result.is_error);
3785 }
3786
3787 #[tokio::test]
3788 async fn test_input_schema_override_wins_over_typed_handler() {
3789 let custom = serde_json::json!({
3793 "type": "object",
3794 "title": "GreetOverride",
3795 "properties": {
3796 "name": { "type": "string", "minLength": 1, "maxLength": 64 }
3797 },
3798 "required": ["name"],
3799 "additionalProperties": false
3800 });
3801
3802 let tool = ToolBuilder::new("greet")
3803 .description("Greet someone with a hand-tuned schema")
3804 .input_schema(custom.clone())
3805 .handler(|input: GreetInput| async move {
3806 Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
3807 })
3808 .build();
3809
3810 let schema = tool.definition().input_schema;
3811 assert_eq!(schema, custom);
3812 assert_eq!(schema["title"], "GreetOverride");
3814
3815 let result = tool.call(serde_json::json!({"name": "World"})).await;
3817 assert!(!result.is_error);
3818 }
3819
3820 #[tokio::test]
3821 async fn test_input_schema_override_preserves_2020_12_constructs() {
3822 let custom = serde_json::json!({
3825 "type": "object",
3826 "properties": {
3827 "filter": {
3828 "oneOf": [
3829 { "type": "string" },
3830 {
3831 "type": "object",
3832 "properties": { "field": { "type": "string" } },
3833 "required": ["field"]
3834 }
3835 ]
3836 }
3837 },
3838 "required": ["filter"]
3839 });
3840
3841 let tool = ToolBuilder::new("filter_tool")
3842 .description("Demonstrates oneOf preservation")
3843 .input_schema(custom.clone())
3844 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3845 Ok(CallToolResult::json(args))
3846 })
3847 .build();
3848
3849 let schema = tool.definition().input_schema;
3850 assert_eq!(schema, custom);
3851 let one_of = schema["properties"]["filter"]["oneOf"]
3852 .as_array()
3853 .expect("oneOf must survive as an array");
3854 assert_eq!(one_of.len(), 2);
3855 assert_eq!(one_of[0]["type"], "string");
3856 assert_eq!(one_of[1]["type"], "object");
3857 }
3858
3859 #[tokio::test]
3860 async fn test_input_schema_override_adds_type_object_if_missing() {
3861 let custom_no_type = serde_json::json!({
3864 "properties": {
3865 "x": { "type": "number" }
3866 }
3867 });
3868
3869 let tool = ToolBuilder::new("typeless")
3870 .description("Schema missing top-level type")
3871 .input_schema(custom_no_type)
3872 .extractor_handler((), |RawArgs(args): RawArgs| async move {
3873 Ok(CallToolResult::json(args))
3874 })
3875 .build();
3876
3877 let schema = tool.definition().input_schema;
3878 assert_eq!(schema["type"], "object");
3879 assert!(schema["properties"]["x"].is_object());
3880 }
3881}