use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use serde::de::DeserializeOwned;
use serde_json::Value;
use crate::error::CapabilityError;
use crate::id::validate_capability_id;
pub use async_trait::async_trait;
pub use schemars::{self, JsonSchema};
pub use serde::{self, Deserialize, Serialize};
pub use serde_json;
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Clone)]
pub struct Definition {
id: String,
name: String,
description: String,
instructions: Option<String>,
metadata: Option<Value>,
tools: Vec<Tool>,
}
impl Definition {
pub fn new(
id: impl Into<String>,
name: impl Into<String>,
description: impl Into<String>,
) -> Self {
Self {
id: id.into(),
name: name.into(),
description: description.into(),
instructions: None,
metadata: None,
tools: Vec::new(),
}
}
pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
self.instructions = Some(instructions.into());
self
}
pub fn metadata(mut self, metadata: impl Into<Value>) -> Self {
self.metadata = Some(metadata.into());
self
}
pub fn tool<H>(mut self, handler: H) -> Self
where
H: Handler,
{
self.tools.push(Tool::new(handler));
self
}
pub fn id(&self) -> &str {
&self.id
}
pub fn name(&self) -> &str {
&self.name
}
pub fn description(&self) -> &str {
&self.description
}
pub fn instructions_text(&self) -> Option<&str> {
self.instructions.as_deref()
}
pub fn metadata_value(&self) -> Option<&Value> {
self.metadata.as_ref()
}
pub fn tools(&self) -> &[Tool] {
&self.tools
}
pub fn validate(&self) -> Result<(), CapabilityError> {
validate_capability_id(&self.id)?;
let invalid = |reason: &str| CapabilityError::InvalidDefinition {
id: self.id.clone(),
reason: reason.to_string(),
};
if self.name.trim().is_empty() {
return Err(invalid("capability name must not be blank"));
}
if self.description.trim().is_empty() {
return Err(invalid("capability description must not be blank"));
}
if self
.instructions
.as_ref()
.is_some_and(|text| text.trim().is_empty())
{
return Err(invalid("capability instructions must not be blank"));
}
if self.tools.is_empty() {
return Err(invalid("capability must define at least one tool"));
}
if let Some(tool) = self
.tools
.iter()
.find(|tool| tool.spec.description.trim().is_empty())
{
return Err(CapabilityError::InvalidDefinition {
id: self.id.clone(),
reason: format!("tool {:?} description must not be blank", tool.spec.name),
});
}
Ok(())
}
}
impl fmt::Debug for Definition {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Definition")
.field("id", &self.id)
.field("name", &self.name)
.field("description", &self.description)
.field("instructions", &self.instructions)
.field("metadata", &self.metadata)
.field("tools", &self.tools)
.finish()
}
}
#[async_trait]
pub trait Handler: Send + Sync + 'static {
type Input: DeserializeOwned + JsonSchema + Send + 'static;
type Output: Serialize + JsonSchema + Send + 'static;
type Error: Into<Error> + Send + 'static;
fn name(&self) -> &str;
fn description(&self) -> &str;
fn display_name(&self) -> Option<&str> {
None
}
fn hints(&self) -> Hints {
Hints::default()
}
async fn execute(
&self,
input: Self::Input,
context: Context,
) -> Result<Self::Output, Self::Error>;
}
#[derive(Clone)]
pub struct Tool {
spec: ToolSpec,
handler: Arc<dyn ErasedHandler>,
}
impl Tool {
fn new<H: Handler>(handler: H) -> Self {
let spec = ToolSpec {
name: handler.name().to_string(),
display_name: handler.display_name().map(str::to_string),
description: handler.description().to_string(),
input_schema: schema_for::<H::Input>(),
output_schema: schema_for::<H::Output>(),
hints: handler.hints(),
};
Self {
spec,
handler: Arc::new(HandlerAdapter(handler)),
}
}
pub fn spec(&self) -> &ToolSpec {
&self.spec
}
pub async fn invoke(&self, arguments: Value, context: Context) -> Result<Value, Error> {
self.handler.call(arguments, context).await
}
}
impl fmt::Debug for Tool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Tool").field("spec", &self.spec).finish()
}
}
#[async_trait]
trait ErasedHandler: Send + Sync {
async fn call(&self, input: Value, context: Context) -> Result<Value, Error>;
}
struct HandlerAdapter<H>(H);
#[async_trait]
impl<H: Handler> ErasedHandler for HandlerAdapter<H> {
async fn call(&self, input: Value, context: Context) -> Result<Value, Error> {
let input = serde_json::from_value::<H::Input>(input).map_err(|error| {
Error::user(
"invalid_arguments",
format!("tool arguments did not match the declared schema: {error}"),
)
})?;
let output = self.0.execute(input, context).await.map_err(Into::into)?;
serde_json::to_value(output).map_err(|error| {
Error::internal(
"result_serialization",
format!("failed to serialize capability result: {error}"),
)
})
}
}
fn schema_for<T: JsonSchema>() -> Value {
serde_json::to_value(schemars::schema_for!(T)).unwrap_or(Value::Null)
}
#[derive(Clone, Debug)]
pub struct ToolSpec {
name: String,
display_name: Option<String>,
description: String,
input_schema: Value,
output_schema: Value,
hints: Hints,
}
impl ToolSpec {
pub fn name(&self) -> &str {
&self.name
}
pub fn display_name(&self) -> Option<&str> {
self.display_name.as_deref()
}
pub fn description(&self) -> &str {
&self.description
}
pub fn input_schema(&self) -> &Value {
&self.input_schema
}
pub fn output_schema(&self) -> &Value {
&self.output_schema
}
pub fn hints(&self) -> &Hints {
&self.hints
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Hints {
pub readonly: Option<bool>,
pub destructive: Option<bool>,
pub idempotent: Option<bool>,
pub open_world: Option<bool>,
pub long_running: Option<bool>,
pub concurrency_class: Option<String>,
pub metadata: Option<Value>,
}
impl Hints {
pub fn readonly(mut self, value: bool) -> Self {
self.readonly = Some(value);
self
}
pub fn destructive(mut self, value: bool) -> Self {
self.destructive = Some(value);
self
}
pub fn idempotent(mut self, value: bool) -> Self {
self.idempotent = Some(value);
self
}
pub fn open_world(mut self, value: bool) -> Self {
self.open_world = Some(value);
self
}
pub fn long_running(mut self, value: bool) -> Self {
self.long_running = Some(value);
self
}
pub fn concurrency_class(mut self, value: impl Into<String>) -> Self {
self.concurrency_class = Some(value.into());
self
}
pub fn metadata(mut self, value: impl Into<Value>) -> Self {
self.metadata = Some(value.into());
self
}
}
#[async_trait]
pub trait ProgressSink: Send + Sync {
async fn emit(&self, tool_name: &str, message: &str);
}
pub trait CancellationSignal: Send + Sync {
fn is_cancelled(&self) -> bool;
fn cancelled<'a>(&'a self) -> BoxFuture<'a, ()>;
}
struct NoopProgressSink;
#[async_trait]
impl ProgressSink for NoopProgressSink {
async fn emit(&self, _tool_name: &str, _message: &str) {}
}
struct NeverCancelled;
impl CancellationSignal for NeverCancelled {
fn is_cancelled(&self) -> bool {
false
}
fn cancelled<'a>(&'a self) -> BoxFuture<'a, ()> {
Box::pin(std::future::pending())
}
}
#[derive(Clone)]
pub struct Context {
tool_name: String,
session_id: String,
workspace_id: String,
locale: Option<String>,
progress: Arc<dyn ProgressSink>,
cancellation: CallCancellation,
}
impl Context {
pub fn new(
tool_name: impl Into<String>,
session_id: impl Into<String>,
workspace_id: impl Into<String>,
) -> Self {
Self {
tool_name: tool_name.into(),
session_id: session_id.into(),
workspace_id: workspace_id.into(),
locale: None,
progress: Arc::new(NoopProgressSink),
cancellation: CallCancellation {
inner: Arc::new(NeverCancelled),
},
}
}
pub fn with_locale(mut self, locale: Option<String>) -> Self {
self.locale = locale;
self
}
pub fn with_progress_sink(mut self, sink: Arc<dyn ProgressSink>) -> Self {
self.progress = sink;
self
}
pub fn with_cancellation_signal(mut self, signal: Arc<dyn CancellationSignal>) -> Self {
self.cancellation = CallCancellation { inner: signal };
self
}
pub fn tool_name(&self) -> &str {
&self.tool_name
}
pub fn session_id(&self) -> &str {
&self.session_id
}
pub fn workspace_id(&self) -> &str {
&self.workspace_id
}
pub fn locale(&self) -> Option<&str> {
self.locale.as_deref()
}
pub fn cancellation(&self) -> &CallCancellation {
&self.cancellation
}
pub async fn progress(&self, message: impl AsRef<str>) {
self.progress.emit(&self.tool_name, message.as_ref()).await;
}
}
impl fmt::Debug for Context {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Context")
.field("tool_name", &self.tool_name)
.field("session_id", &self.session_id)
.field("workspace_id", &self.workspace_id)
.field("locale", &self.locale)
.field("cancelled", &self.cancellation.is_cancelled())
.finish()
}
}
#[derive(Clone)]
pub struct CallCancellation {
inner: Arc<dyn CancellationSignal>,
}
impl CallCancellation {
pub fn is_cancelled(&self) -> bool {
self.inner.is_cancelled()
}
pub async fn cancelled(&self) {
self.inner.cancelled().await;
}
}
impl fmt::Debug for CallCancellation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CallCancellation")
.field("cancelled", &self.is_cancelled())
.finish()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum ErrorVisibility {
User,
Internal,
}
#[derive(Debug)]
pub struct Error {
visibility: ErrorVisibility,
code: String,
message: String,
details: Option<Value>,
}
impl Error {
pub fn user(code: impl Into<String>, message: impl Into<String>) -> Self {
Self {
visibility: ErrorVisibility::User,
code: code.into(),
message: message.into(),
details: None,
}
}
pub fn internal(code: impl Into<String>, message: impl Into<String>) -> Self {
Self {
visibility: ErrorVisibility::Internal,
code: code.into(),
message: message.into(),
details: None,
}
}
pub fn details(mut self, details: impl Into<Value>) -> Self {
self.details = Some(details.into());
self
}
pub fn code(&self) -> &str {
&self.code
}
pub fn message(&self) -> &str {
&self.message
}
pub fn details_value(&self) -> Option<&Value> {
self.details.as_ref()
}
pub fn visibility(&self) -> ErrorVisibility {
self.visibility
}
pub fn into_parts(self) -> (ErrorVisibility, String, String, Option<Value>) {
(self.visibility, self.code, self.message, self.details)
}
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "[{}] {}", self.code, self.message)
}
}
impl std::error::Error for Error {}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[derive(Deserialize, JsonSchema)]
struct LookupInput {
city: String,
}
#[derive(Debug, Serialize, JsonSchema)]
struct LookupOutput {
city: String,
temperatures: Vec<i32>,
}
struct Lookup;
#[async_trait]
impl Handler for Lookup {
type Input = LookupInput;
type Output = LookupOutput;
type Error = Error;
fn name(&self) -> &str {
"lookup_weather"
}
fn description(&self) -> &str {
"Look up a typed weather forecast."
}
fn hints(&self) -> Hints {
Hints::default().readonly(true).idempotent(true)
}
async fn execute(
&self,
input: Self::Input,
context: Context,
) -> Result<Self::Output, Self::Error> {
context.progress("forecast ready").await;
Ok(LookupOutput {
city: input.city,
temperatures: vec![18, 21],
})
}
}
fn lookup_capability() -> Definition {
Definition::new("weather", "Weather", "Typed weather tools.")
.metadata(json!({ "owner": "example" }))
.tool(Lookup)
}
#[test]
fn exposes_input_output_schemas_and_hints() {
let capability = lookup_capability();
let spec = capability.tools()[0].spec();
assert_eq!(spec.name(), "lookup_weather");
assert_eq!(spec.input_schema()["type"], "object");
assert_eq!(spec.output_schema()["type"], "object");
assert_eq!(spec.hints().readonly, Some(true));
assert_eq!(
capability.metadata_value(),
Some(&json!({ "owner": "example" }))
);
capability.validate().unwrap();
}
#[test]
fn validate_rejects_structural_problems() {
let err = Definition::new("weather", "Weather", "d")
.validate()
.unwrap_err();
assert!(err.reason().contains("at least one tool"));
let err = Definition::new("2fast", "N", "d")
.tool(Lookup)
.validate()
.unwrap_err();
assert!(err.reason().contains("start with a letter"));
let err = Definition::new("weather", " ", "d")
.tool(Lookup)
.validate()
.unwrap_err();
assert!(err.reason().contains("name must not be blank"));
}
fn block_on<F: Future>(mut future: F) -> F::Output {
use std::task::{Context as TaskContext, Poll, RawWaker, RawWakerVTable, Waker};
fn raw_waker() -> RawWaker {
fn no_op(_: *const ()) {}
fn clone(_: *const ()) -> RawWaker {
raw_waker()
}
RawWaker::new(
std::ptr::null(),
&RawWakerVTable::new(clone, no_op, no_op, no_op),
)
}
let waker = unsafe { Waker::from_raw(raw_waker()) };
let mut context = TaskContext::from_waker(&waker);
let mut future = unsafe { Pin::new_unchecked(&mut future) };
loop {
match future.as_mut().poll(&mut context) {
Poll::Ready(output) => return output,
Poll::Pending => std::thread::yield_now(),
}
}
}
#[test]
fn invoke_serializes_typed_output() {
let tool = lookup_capability().tools()[0].clone();
let context = Context::new("lookup_weather", "session", "workspace");
let value = block_on(tool.invoke(json!({ "city": "Kyiv" }), context)).unwrap();
assert_eq!(value, json!({ "city": "Kyiv", "temperatures": [18, 21] }));
}
#[test]
fn invoke_rejects_invalid_arguments_as_user_error() {
let tool = lookup_capability().tools()[0].clone();
let context = Context::new("lookup_weather", "session", "workspace");
let error = block_on(tool.invoke(json!({ "city": 42 }), context)).unwrap_err();
assert_eq!(error.code(), "invalid_arguments");
assert_eq!(error.visibility(), ErrorVisibility::User);
}
#[test]
fn context_defaults_are_inert() {
let context = Context::new("t", "s", "w");
assert!(!context.cancellation().is_cancelled());
block_on(context.progress("no-op"));
let debug = format!("{context:?}");
assert!(debug.contains("cancelled: false"));
}
}