mod png;
mod session;
mod tree;
use std::sync::Arc;
use aither_core::llm::tool::Tools;
use aither_mcp::transport::{BidirectionalTransport, StdioTransport};
use aither_mcp::{McpError, McpServer};
use waterui_mcp_protocol::{INSTRUCTIONS, register_session_tools};
use waterui_testing::OffscreenApp;
pub use aither_mcp::protocol::ServerInfo;
use session::{Session, SessionHandle};
pub use waterui_mcp_protocol::SESSION_TOOL_NAMES;
pub fn serve_stdio(mount: impl FnMut() -> OffscreenApp, info: ServerInfo) -> Result<(), McpError> {
serve(StdioTransport::new(), mount, info)
}
pub fn serve<T>(
transport: T,
mount: impl FnMut() -> OffscreenApp,
info: ServerInfo,
) -> Result<(), McpError>
where
T: BidirectionalTransport + Sync + Send + 'static,
{
let mut session = Session::new(mount);
let (tx, rx) = async_channel::unbounded();
let mut tools = Tools::new();
{
let handle = Arc::new(SessionHandle::new(tx));
register_session_tools(&mut tools, handle);
debug_assert_eq!(tools.definitions().len(), SESSION_TOOL_NAMES.len());
}
let ServerInfo { name, version } = info;
let server = std::thread::Builder::new()
.name("waterui-mcp-server".to_owned())
.spawn(move || {
futures_lite::future::block_on(
McpServer::new(transport, tools, name, version.unwrap_or_default())
.with_instructions(INSTRUCTIONS)
.run(),
)
})
.map_err(|error| {
McpError::Transport(format!(
"failed to spawn waterui-mcp server thread: {error}"
))
})?;
while let Ok(command) = futures_lite::future::block_on(rx.recv()) {
session.execute(command);
}
match server.join() {
Ok(result) => result,
Err(payload) => std::panic::resume_unwind(payload),
}
}