#![allow(dead_code)]
use liwe::model::config::Configuration;
use rmcp::model::*;
use rmcp::{ClientHandler, ServiceExt, service::RunningService, RoleClient};
use iwec::IweServer;
#[derive(Default, Clone)]
struct TestClient;
impl ClientHandler for TestClient {}
pub struct Fixture {
client: RunningService<RoleClient, TestClient>,
_server_handle: tokio::task::JoinHandle<anyhow::Result<()>>,
}
impl Fixture {
pub async fn with_documents(documents: Vec<(&str, &str)>) -> Self {
Self::with_documents_and_config(documents, Configuration::default()).await
}
pub async fn with_documents_and_config(documents: Vec<(&str, &str)>, config: Configuration) -> Self {
let server = IweServer::from_documents_with_config(documents, config);
let (server_transport, client_transport) = tokio::io::duplex(65536);
let server_handle = tokio::spawn(async move {
let service = server.serve(server_transport).await?;
service.waiting().await?;
anyhow::Ok(())
});
let client = TestClient.serve(client_transport).await.expect("client to connect");
Self {
client,
_server_handle: server_handle,
}
}
pub async fn call_tool(&self, name: &str, arguments: serde_json::Value) -> CallToolResult {
self.try_call_tool(name, arguments)
.await
.expect("tool call to succeed")
}
pub async fn try_call_tool(
&self,
name: &str,
arguments: serde_json::Value,
) -> Result<CallToolResult, rmcp::ServiceError> {
let args = match arguments {
serde_json::Value::Object(map) => map,
serde_json::Value::Null => serde_json::Map::new(),
other => panic!("arguments must be an object, got: {other}"),
};
let params = CallToolRequestParams::new(name.to_string()).with_arguments(args);
let response = self
.client
.send_request(ClientRequest::CallToolRequest(Request::new(params)))
.await?;
match response {
ServerResult::CallToolResult(result) => Ok(result),
other => panic!("expected CallToolResult, got: {other:?}"),
}
}
pub fn result_text(result: &CallToolResult) -> String {
result
.content
.iter()
.filter_map(|c| match &c.raw {
RawContent::Text(t) => Some(t.text.clone()),
_ => None,
})
.collect::<Vec<_>>()
.join("")
}
pub fn result_json(result: &CallToolResult) -> serde_json::Value {
let text = Self::result_text(result);
serde_json::from_str(&text).expect("result to be valid JSON")
}
pub async fn list_resources(&self) -> ListResourcesResult {
let response = self
.client
.send_request(ClientRequest::ListResourcesRequest(
rmcp::model::RequestOptionalParam::default(),
))
.await
.expect("list resources to succeed");
match response {
ServerResult::ListResourcesResult(result) => result,
other => panic!("expected ListResourcesResult, got: {other:?}"),
}
}
pub async fn read_resource(&self, uri: &str) -> ReadResourceResult {
let params = ReadResourceRequestParams::new(uri.to_string());
let response = self
.client
.send_request(ClientRequest::ReadResourceRequest(Request::new(params)))
.await
.expect("read resource to succeed");
match response {
ServerResult::ReadResourceResult(result) => result,
other => panic!("expected ReadResourceResult, got: {other:?}"),
}
}
pub async fn list_prompts(&self) -> ListPromptsResult {
let response = self
.client
.send_request(ClientRequest::ListPromptsRequest(
rmcp::model::RequestOptionalParam::default(),
))
.await
.expect("list prompts to succeed");
match response {
ServerResult::ListPromptsResult(result) => result,
other => panic!("expected ListPromptsResult, got: {other:?}"),
}
}
pub async fn get_prompt(
&self,
name: &str,
arguments: serde_json::Value,
) -> GetPromptResult {
let params = match arguments {
serde_json::Value::Object(map) => {
let string_map: serde_json::Map<String, serde_json::Value> = map
.into_iter()
.map(|(k, v)| {
let s = v.as_str().unwrap_or_default().to_string();
(k, serde_json::Value::String(s))
})
.collect();
GetPromptRequestParams::new(name.to_string()).with_arguments(string_map)
}
serde_json::Value::Null => GetPromptRequestParams::new(name.to_string()),
other => panic!("arguments must be an object, got: {other}"),
};
let response = self
.client
.send_request(ClientRequest::GetPromptRequest(Request::new(params)))
.await
.expect("get prompt to succeed");
match response {
ServerResult::GetPromptResult(result) => result,
other => panic!("expected GetPromptResult, got: {other:?}"),
}
}
}