1use std::fmt;
67use std::future::Future;
68use std::pin::Pin;
69use std::sync::Arc;
70
71use serde::de::DeserializeOwned;
72use serde_json::Value;
73
74use crate::error::CapabilityError;
75use crate::id::validate_capability_id;
76
77pub use async_trait::async_trait;
78pub use schemars::{self, JsonSchema};
79pub use serde::{self, Deserialize, Serialize};
80pub use serde_json;
81
82pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
84
85#[derive(Clone)]
92pub struct Definition {
93 id: String,
94 name: String,
95 description: String,
96 instructions: Option<String>,
97 metadata: Option<Value>,
98 tools: Vec<Tool>,
99}
100
101impl Definition {
102 pub fn new(
109 id: impl Into<String>,
110 name: impl Into<String>,
111 description: impl Into<String>,
112 ) -> Self {
113 Self {
114 id: id.into(),
115 name: name.into(),
116 description: description.into(),
117 instructions: None,
118 metadata: None,
119 tools: Vec::new(),
120 }
121 }
122
123 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
128 self.instructions = Some(instructions.into());
129 self
130 }
131
132 pub fn metadata(mut self, metadata: impl Into<Value>) -> Self {
137 self.metadata = Some(metadata.into());
138 self
139 }
140
141 pub fn tool<H>(mut self, handler: H) -> Self
143 where
144 H: Handler,
145 {
146 self.tools.push(Tool::new(handler));
147 self
148 }
149
150 pub fn id(&self) -> &str {
152 &self.id
153 }
154
155 pub fn name(&self) -> &str {
157 &self.name
158 }
159
160 pub fn description(&self) -> &str {
162 &self.description
163 }
164
165 pub fn instructions_text(&self) -> Option<&str> {
167 self.instructions.as_deref()
168 }
169
170 pub fn metadata_value(&self) -> Option<&Value> {
172 self.metadata.as_ref()
173 }
174
175 pub fn tools(&self) -> &[Tool] {
177 &self.tools
178 }
179
180 pub fn validate(&self) -> Result<(), CapabilityError> {
182 validate_capability_id(&self.id)?;
183 let invalid = |reason: &str| CapabilityError::InvalidDefinition {
184 id: self.id.clone(),
185 reason: reason.to_string(),
186 };
187 if self.name.trim().is_empty() {
188 return Err(invalid("capability name must not be blank"));
189 }
190 if self.description.trim().is_empty() {
191 return Err(invalid("capability description must not be blank"));
192 }
193 if self
194 .instructions
195 .as_ref()
196 .is_some_and(|text| text.trim().is_empty())
197 {
198 return Err(invalid("capability instructions must not be blank"));
199 }
200 if self.tools.is_empty() {
201 return Err(invalid("capability must define at least one tool"));
202 }
203 if let Some(tool) = self
204 .tools
205 .iter()
206 .find(|tool| tool.spec.description.trim().is_empty())
207 {
208 return Err(CapabilityError::InvalidDefinition {
209 id: self.id.clone(),
210 reason: format!("tool {:?} description must not be blank", tool.spec.name),
211 });
212 }
213 Ok(())
214 }
215}
216
217impl fmt::Debug for Definition {
218 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
219 f.debug_struct("Definition")
220 .field("id", &self.id)
221 .field("name", &self.name)
222 .field("description", &self.description)
223 .field("instructions", &self.instructions)
224 .field("metadata", &self.metadata)
225 .field("tools", &self.tools)
226 .finish()
227 }
228}
229
230#[async_trait]
238pub trait Handler: Send + Sync + 'static {
239 type Input: DeserializeOwned + JsonSchema + Send + 'static;
241 type Output: Serialize + JsonSchema + Send + 'static;
243 type Error: Into<Error> + Send + 'static;
245
246 fn name(&self) -> &str;
248 fn description(&self) -> &str;
250
251 fn display_name(&self) -> Option<&str> {
253 None
254 }
255
256 fn hints(&self) -> Hints {
258 Hints::default()
259 }
260
261 async fn execute(
267 &self,
268 input: Self::Input,
269 context: Context,
270 ) -> Result<Self::Output, Self::Error>;
271}
272
273#[derive(Clone)]
279pub struct Tool {
280 spec: ToolSpec,
281 handler: Arc<dyn ErasedHandler>,
282}
283
284impl Tool {
285 fn new<H: Handler>(handler: H) -> Self {
286 let spec = ToolSpec {
287 name: handler.name().to_string(),
288 display_name: handler.display_name().map(str::to_string),
289 description: handler.description().to_string(),
290 input_schema: json_schema_for::<H::Input>(),
291 output_schema: json_schema_for::<H::Output>(),
292 hints: handler.hints(),
293 };
294 Self {
295 spec,
296 handler: Arc::new(HandlerAdapter(handler)),
297 }
298 }
299
300 pub fn spec(&self) -> &ToolSpec {
302 &self.spec
303 }
304
305 pub async fn invoke(&self, arguments: Value, context: Context) -> Result<Value, Error> {
312 self.handler.call(arguments, context).await
313 }
314}
315
316impl fmt::Debug for Tool {
317 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
318 f.debug_struct("Tool").field("spec", &self.spec).finish()
319 }
320}
321
322#[async_trait]
323trait ErasedHandler: Send + Sync {
324 async fn call(&self, input: Value, context: Context) -> Result<Value, Error>;
325}
326
327struct HandlerAdapter<H>(H);
328
329#[async_trait]
330impl<H: Handler> ErasedHandler for HandlerAdapter<H> {
331 async fn call(&self, input: Value, context: Context) -> Result<Value, Error> {
332 let input = serde_json::from_value::<H::Input>(input).map_err(|error| {
333 Error::user(
334 "invalid_arguments",
335 format!("tool arguments did not match the declared schema: {error}"),
336 )
337 })?;
338 let output = self.0.execute(input, context).await.map_err(Into::into)?;
339 serde_json::to_value(output).map_err(|error| {
340 Error::internal(
341 "result_serialization",
342 format!("failed to serialize capability result: {error}"),
343 )
344 })
345 }
346}
347
348pub fn json_schema_for<T: JsonSchema>() -> Value {
390 serde_json::to_value(schemars::schema_for!(T)).unwrap_or(Value::Null)
391}
392
393#[derive(Clone, Debug)]
395pub struct ToolSpec {
396 name: String,
397 display_name: Option<String>,
398 description: String,
399 input_schema: Value,
400 output_schema: Value,
401 hints: Hints,
402}
403
404impl ToolSpec {
405 pub fn name(&self) -> &str {
407 &self.name
408 }
409 pub fn display_name(&self) -> Option<&str> {
411 self.display_name.as_deref()
412 }
413 pub fn description(&self) -> &str {
415 &self.description
416 }
417 pub fn input_schema(&self) -> &Value {
419 &self.input_schema
420 }
421 pub fn output_schema(&self) -> &Value {
423 &self.output_schema
424 }
425 pub fn hints(&self) -> &Hints {
427 &self.hints
428 }
429}
430
431#[derive(Clone, Debug, Default, PartialEq)]
436pub struct Hints {
437 pub readonly: Option<bool>,
439 pub destructive: Option<bool>,
441 pub idempotent: Option<bool>,
443 pub open_world: Option<bool>,
445 pub long_running: Option<bool>,
447 pub concurrency_class: Option<String>,
449 pub metadata: Option<Value>,
451}
452
453impl Hints {
454 pub fn readonly(mut self, value: bool) -> Self {
456 self.readonly = Some(value);
457 self
458 }
459 pub fn destructive(mut self, value: bool) -> Self {
461 self.destructive = Some(value);
462 self
463 }
464 pub fn idempotent(mut self, value: bool) -> Self {
466 self.idempotent = Some(value);
467 self
468 }
469 pub fn open_world(mut self, value: bool) -> Self {
471 self.open_world = Some(value);
472 self
473 }
474 pub fn long_running(mut self, value: bool) -> Self {
476 self.long_running = Some(value);
477 self
478 }
479 pub fn concurrency_class(mut self, value: impl Into<String>) -> Self {
481 self.concurrency_class = Some(value.into());
482 self
483 }
484 pub fn metadata(mut self, value: impl Into<Value>) -> Self {
486 self.metadata = Some(value.into());
487 self
488 }
489}
490
491#[async_trait]
496pub trait ProgressSink: Send + Sync {
497 async fn emit(&self, tool_name: &str, message: &str);
499}
500
501pub trait CancellationSignal: Send + Sync {
505 fn is_cancelled(&self) -> bool;
507
508 fn cancelled<'a>(&'a self) -> BoxFuture<'a, ()>;
510}
511
512struct NoopProgressSink;
513
514#[async_trait]
515impl ProgressSink for NoopProgressSink {
516 async fn emit(&self, _tool_name: &str, _message: &str) {}
517}
518
519struct NeverCancelled;
520
521impl CancellationSignal for NeverCancelled {
522 fn is_cancelled(&self) -> bool {
523 false
524 }
525
526 fn cancelled<'a>(&'a self) -> BoxFuture<'a, ()> {
527 Box::pin(std::future::pending())
528 }
529}
530
531#[derive(Clone)]
539pub struct Context {
540 tool_name: String,
541 session_id: String,
542 workspace_id: String,
543 locale: Option<String>,
544 progress: Arc<dyn ProgressSink>,
545 cancellation: CallCancellation,
546}
547
548impl Context {
549 pub fn new(
551 tool_name: impl Into<String>,
552 session_id: impl Into<String>,
553 workspace_id: impl Into<String>,
554 ) -> Self {
555 Self {
556 tool_name: tool_name.into(),
557 session_id: session_id.into(),
558 workspace_id: workspace_id.into(),
559 locale: None,
560 progress: Arc::new(NoopProgressSink),
561 cancellation: CallCancellation {
562 inner: Arc::new(NeverCancelled),
563 },
564 }
565 }
566
567 pub fn with_locale(mut self, locale: Option<String>) -> Self {
569 self.locale = locale;
570 self
571 }
572
573 pub fn with_progress_sink(mut self, sink: Arc<dyn ProgressSink>) -> Self {
575 self.progress = sink;
576 self
577 }
578
579 pub fn with_cancellation_signal(mut self, signal: Arc<dyn CancellationSignal>) -> Self {
581 self.cancellation = CallCancellation { inner: signal };
582 self
583 }
584
585 pub fn tool_name(&self) -> &str {
587 &self.tool_name
588 }
589
590 pub fn session_id(&self) -> &str {
592 &self.session_id
593 }
594
595 pub fn workspace_id(&self) -> &str {
597 &self.workspace_id
598 }
599
600 pub fn locale(&self) -> Option<&str> {
602 self.locale.as_deref()
603 }
604
605 pub fn cancellation(&self) -> &CallCancellation {
607 &self.cancellation
608 }
609
610 pub async fn progress(&self, message: impl AsRef<str>) {
614 self.progress.emit(&self.tool_name, message.as_ref()).await;
615 }
616}
617
618impl fmt::Debug for Context {
619 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
620 f.debug_struct("Context")
621 .field("tool_name", &self.tool_name)
622 .field("session_id", &self.session_id)
623 .field("workspace_id", &self.workspace_id)
624 .field("locale", &self.locale)
625 .field("cancelled", &self.cancellation.is_cancelled())
626 .finish()
627 }
628}
629
630#[derive(Clone)]
637pub struct CallCancellation {
638 inner: Arc<dyn CancellationSignal>,
639}
640
641impl CallCancellation {
642 pub fn is_cancelled(&self) -> bool {
644 self.inner.is_cancelled()
645 }
646
647 pub async fn cancelled(&self) {
649 self.inner.cancelled().await;
650 }
651}
652
653impl fmt::Debug for CallCancellation {
654 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
655 f.debug_struct("CallCancellation")
656 .field("cancelled", &self.is_cancelled())
657 .finish()
658 }
659}
660
661#[derive(Clone, Copy, Debug, PartialEq, Eq)]
663#[non_exhaustive]
664pub enum ErrorVisibility {
665 User,
667 Internal,
669}
670
671#[derive(Debug)]
679pub struct Error {
680 visibility: ErrorVisibility,
681 code: String,
682 message: String,
683 details: Option<Value>,
684}
685
686impl Error {
687 pub fn user(code: impl Into<String>, message: impl Into<String>) -> Self {
689 Self {
690 visibility: ErrorVisibility::User,
691 code: code.into(),
692 message: message.into(),
693 details: None,
694 }
695 }
696
697 pub fn internal(code: impl Into<String>, message: impl Into<String>) -> Self {
699 Self {
700 visibility: ErrorVisibility::Internal,
701 code: code.into(),
702 message: message.into(),
703 details: None,
704 }
705 }
706
707 pub fn details(mut self, details: impl Into<Value>) -> Self {
709 self.details = Some(details.into());
710 self
711 }
712
713 pub fn code(&self) -> &str {
715 &self.code
716 }
717
718 pub fn message(&self) -> &str {
720 &self.message
721 }
722
723 pub fn details_value(&self) -> Option<&Value> {
725 self.details.as_ref()
726 }
727
728 pub fn visibility(&self) -> ErrorVisibility {
730 self.visibility
731 }
732
733 pub fn into_parts(self) -> (ErrorVisibility, String, String, Option<Value>) {
735 (self.visibility, self.code, self.message, self.details)
736 }
737}
738
739impl fmt::Display for Error {
740 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
741 write!(f, "[{}] {}", self.code, self.message)
742 }
743}
744
745impl std::error::Error for Error {}
746
747#[cfg(test)]
748mod tests {
749 use super::*;
750 use serde_json::json;
751
752 #[derive(Deserialize, JsonSchema)]
753 struct LookupInput {
754 city: String,
755 }
756
757 #[derive(Debug, Serialize, JsonSchema)]
758 struct LookupOutput {
759 city: String,
760 temperatures: Vec<i32>,
761 }
762
763 struct Lookup;
764
765 #[async_trait]
766 impl Handler for Lookup {
767 type Input = LookupInput;
768 type Output = LookupOutput;
769 type Error = Error;
770
771 fn name(&self) -> &str {
772 "lookup_weather"
773 }
774
775 fn description(&self) -> &str {
776 "Look up a typed weather forecast."
777 }
778
779 fn hints(&self) -> Hints {
780 Hints::default().readonly(true).idempotent(true)
781 }
782
783 async fn execute(
784 &self,
785 input: Self::Input,
786 context: Context,
787 ) -> Result<Self::Output, Self::Error> {
788 context.progress("forecast ready").await;
789 Ok(LookupOutput {
790 city: input.city,
791 temperatures: vec![18, 21],
792 })
793 }
794 }
795
796 fn lookup_capability() -> Definition {
797 Definition::new("weather", "Weather", "Typed weather tools.")
798 .metadata(json!({ "owner": "example" }))
799 .tool(Lookup)
800 }
801
802 #[test]
803 fn exposes_input_output_schemas_and_hints() {
804 let capability = lookup_capability();
805 let spec = capability.tools()[0].spec();
806
807 assert_eq!(spec.name(), "lookup_weather");
808 assert_eq!(spec.input_schema()["type"], "object");
809 assert_eq!(spec.input_schema()["required"], json!(["city"]));
810 assert_eq!(spec.input_schema()["properties"]["city"]["type"], "string");
811 assert_eq!(spec.output_schema()["type"], "object");
812 assert_eq!(spec.output_schema()["properties"]["city"]["type"], "string");
813 assert_eq!(
814 spec.output_schema()["properties"]["temperatures"]["type"],
815 "array"
816 );
817 assert_eq!(
818 spec.output_schema()["properties"]["temperatures"]["items"]["type"],
819 "integer"
820 );
821 assert_eq!(spec.hints().readonly, Some(true));
822 assert_eq!(spec.hints().idempotent, Some(true));
823 assert_eq!(
824 capability.metadata_value(),
825 Some(&json!({ "owner": "example" }))
826 );
827 capability.validate().unwrap();
828 }
829
830 #[test]
831 fn validate_rejects_structural_problems() {
832 let err = Definition::new("weather", "Weather", "d")
833 .validate()
834 .unwrap_err();
835 assert!(err.reason().contains("at least one tool"));
836
837 let err = Definition::new("2fast", "N", "d")
838 .tool(Lookup)
839 .validate()
840 .unwrap_err();
841 assert!(err.reason().contains("start with a letter"));
842
843 let err = Definition::new("weather", " ", "d")
844 .tool(Lookup)
845 .validate()
846 .unwrap_err();
847 assert!(err.reason().contains("name must not be blank"));
848 for (definition, expected) in [
849 (
850 Definition::new("weather", "Weather", " ").tool(Lookup),
851 "description must not be blank",
852 ),
853 (
854 lookup_capability().instructions(" "),
855 "instructions must not be blank",
856 ),
857 ] {
858 assert!(
859 definition
860 .validate()
861 .unwrap_err()
862 .reason()
863 .contains(expected)
864 );
865 }
866 let mut malformed = lookup_capability();
867 malformed.tools[0].spec.description = " ".into();
868 assert!(
869 malformed
870 .validate()
871 .unwrap_err()
872 .reason()
873 .contains("tool \"lookup_weather\" description must not be blank")
874 );
875 }
876
877 fn ready<F: Future>(future: F) -> F::Output {
879 let mut future = std::pin::pin!(future);
880 let mut context = std::task::Context::from_waker(std::task::Waker::noop());
881 match future.as_mut().poll(&mut context) {
882 std::task::Poll::Ready(value) => value,
883 std::task::Poll::Pending => panic!("test future unexpectedly pending"),
884 }
885 }
886
887 #[derive(Default)]
888 struct RecordingProgress(std::sync::Mutex<Vec<(String, String)>>);
889 #[async_trait]
890 impl ProgressSink for RecordingProgress {
891 async fn emit(&self, tool: &str, message: &str) {
892 self.0.lock().unwrap().push((tool.into(), message.into()));
893 }
894 }
895
896 #[test]
897 fn invoke_serializes_typed_output() {
898 let tool = lookup_capability().tools()[0].clone();
899 let progress = Arc::new(RecordingProgress::default());
900 let context = Context::new("lookup_weather", "session", "workspace")
901 .with_progress_sink(progress.clone());
902 let value = ready(tool.invoke(json!({ "city": "Kyiv" }), context)).unwrap();
903 assert_eq!(value, json!({ "city": "Kyiv", "temperatures": [18, 21] }));
904 assert_eq!(
905 *progress.0.lock().unwrap(),
906 [("lookup_weather".into(), "forecast ready".into())]
907 );
908 }
909
910 #[test]
911 fn invoke_rejects_invalid_arguments_as_user_error() {
912 let tool = lookup_capability().tools()[0].clone();
913 let context = Context::new("lookup_weather", "session", "workspace");
914 let error = ready(tool.invoke(json!({ "city": 42 }), context)).unwrap_err();
915 assert_eq!(error.code(), "invalid_arguments");
916 assert_eq!(error.visibility(), ErrorVisibility::User);
917 }
918
919 #[test]
920 fn context_defaults_are_inert() {
921 let context = Context::new("t", "s", "w");
922 assert!(!context.cancellation().is_cancelled());
923 ready(context.progress("no-op"));
924 let mut cancelled = std::pin::pin!(context.cancellation().cancelled());
925 let mut task = std::task::Context::from_waker(std::task::Waker::noop());
926 assert!(cancelled.as_mut().poll(&mut task).is_pending());
927 }
928}