use std::{future::Future, pin::Pin};
use crate::{
ToolError,
execution::{ToolHandlerInput, ToolHandlerOutput},
sources::{local, mcp},
};
pub trait ToolHandlerBackend: Send + 'static {
fn step(
&mut self,
input: ToolHandlerInput,
) -> impl Future<Output = Result<ToolHandlerOutput, ToolError>> + Send + '_;
}
pub struct ToolHandler {
inner: InnerToolHandler,
}
enum InnerToolHandler {
Local(local::handler::LocalToolHandler),
Mcp(mcp::handler::McpToolHandler),
Extension(Box<dyn ErasedToolHandlerBackend>),
}
impl ToolHandler {
pub(crate) fn from_local(handler: local::handler::LocalToolHandler) -> Self {
Self {
inner: InnerToolHandler::Local(handler),
}
}
pub(crate) fn from_mcp(handler: mcp::handler::McpToolHandler) -> Self {
Self {
inner: InnerToolHandler::Mcp(handler),
}
}
pub(crate) fn from_extension<B: ToolHandlerBackend>(backend: B) -> Self {
Self {
inner: InnerToolHandler::Extension(Box::new(backend)),
}
}
}
impl std::fmt::Debug for ToolHandler {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("ToolHandler")
.finish_non_exhaustive()
}
}
trait ErasedToolHandlerBackend: Send {
fn step_erased(
&mut self,
input: ToolHandlerInput,
) -> Pin<Box<dyn Future<Output = Result<ToolHandlerOutput, ToolError>> + Send + '_>>;
}
impl<B: ToolHandlerBackend> ErasedToolHandlerBackend for B {
fn step_erased(
&mut self,
input: ToolHandlerInput,
) -> Pin<Box<dyn Future<Output = Result<ToolHandlerOutput, ToolError>> + Send + '_>> {
Box::pin(self.step(input))
}
}
impl ToolHandlerBackend for InnerToolHandler {
async fn step(&mut self, input: ToolHandlerInput) -> Result<ToolHandlerOutput, ToolError> {
match self {
Self::Local(handler) => handler.step(input).await,
Self::Mcp(handler) => handler.step(input).await,
Self::Extension(handler) => handler.step_erased(input).await,
}
}
}
impl ToolHandlerBackend for ToolHandler {
fn step(
&mut self,
input: ToolHandlerInput,
) -> impl Future<Output = Result<ToolHandlerOutput, ToolError>> + Send + '_ {
self.inner.step(input)
}
}