context server: Make requests type safe (#32254)
This changes the context server crate so that the input/output for a request are encoded at the type level, similar to how it is done for LSP requests. This also makes it easier to write tests that mock context servers, e.g. you can write something like this now when using the `test-support` feature of the `context-server` crate: ```rust create_fake_transport("mcp-1", cx.background_executor()) .on_request::<context_server::types::request::PromptsList>(|_params| { PromptsListResponse { prompts: vec![/* some prompts */], .. } }) ``` Release Notes: - N/A
This commit is contained in:
parent
454adfacae
commit
95d78ff8d5
11 changed files with 320 additions and 433 deletions
|
@ -6,10 +6,9 @@
|
|||
//! of messages.
|
||||
|
||||
use anyhow::Result;
|
||||
use collections::HashMap;
|
||||
|
||||
use crate::client::Client;
|
||||
use crate::types;
|
||||
use crate::types::{self, Request};
|
||||
|
||||
pub struct ModelContextProtocol {
|
||||
inner: Client,
|
||||
|
@ -43,7 +42,7 @@ impl ModelContextProtocol {
|
|||
|
||||
let response: types::InitializeResponse = self
|
||||
.inner
|
||||
.request(types::RequestType::Initialize.as_str(), params)
|
||||
.request(types::request::Initialize::METHOD, params)
|
||||
.await?;
|
||||
|
||||
anyhow::ensure!(
|
||||
|
@ -94,137 +93,7 @@ impl InitializedContextServerProtocol {
|
|||
}
|
||||
}
|
||||
|
||||
fn check_capability(&self, capability: ServerCapability) -> Result<()> {
|
||||
anyhow::ensure!(
|
||||
self.capable(capability),
|
||||
"Server does not support {capability:?} capability"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// List the MCP prompts.
|
||||
pub async fn list_prompts(&self) -> Result<Vec<types::Prompt>> {
|
||||
self.check_capability(ServerCapability::Prompts)?;
|
||||
|
||||
let response: types::PromptsListResponse = self
|
||||
.inner
|
||||
.request(
|
||||
types::RequestType::PromptsList.as_str(),
|
||||
serde_json::json!({}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(response.prompts)
|
||||
}
|
||||
|
||||
/// List the MCP resources.
|
||||
pub async fn list_resources(&self) -> Result<types::ResourcesListResponse> {
|
||||
self.check_capability(ServerCapability::Resources)?;
|
||||
|
||||
let response: types::ResourcesListResponse = self
|
||||
.inner
|
||||
.request(
|
||||
types::RequestType::ResourcesList.as_str(),
|
||||
serde_json::json!({}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Executes a prompt with the given arguments and returns the result.
|
||||
pub async fn run_prompt<P: AsRef<str>>(
|
||||
&self,
|
||||
prompt: P,
|
||||
arguments: HashMap<String, String>,
|
||||
) -> Result<types::PromptsGetResponse> {
|
||||
self.check_capability(ServerCapability::Prompts)?;
|
||||
|
||||
let params = types::PromptsGetParams {
|
||||
name: prompt.as_ref().to_string(),
|
||||
arguments: Some(arguments),
|
||||
meta: None,
|
||||
};
|
||||
|
||||
let response: types::PromptsGetResponse = self
|
||||
.inner
|
||||
.request(types::RequestType::PromptsGet.as_str(), params)
|
||||
.await?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
pub async fn completion<P: Into<String>>(
|
||||
&self,
|
||||
reference: types::CompletionReference,
|
||||
argument: P,
|
||||
value: P,
|
||||
) -> Result<types::Completion> {
|
||||
let params = types::CompletionCompleteParams {
|
||||
r#ref: reference,
|
||||
argument: types::CompletionArgument {
|
||||
name: argument.into(),
|
||||
value: value.into(),
|
||||
},
|
||||
meta: None,
|
||||
};
|
||||
let result: types::CompletionCompleteResponse = self
|
||||
.inner
|
||||
.request(types::RequestType::CompletionComplete.as_str(), params)
|
||||
.await?;
|
||||
|
||||
let completion = types::Completion {
|
||||
values: result.completion.values,
|
||||
total: types::CompletionTotal::from_options(
|
||||
result.completion.has_more,
|
||||
result.completion.total,
|
||||
),
|
||||
};
|
||||
|
||||
Ok(completion)
|
||||
}
|
||||
|
||||
/// List MCP tools.
|
||||
pub async fn list_tools(&self) -> Result<types::ListToolsResponse> {
|
||||
self.check_capability(ServerCapability::Tools)?;
|
||||
|
||||
let response = self
|
||||
.inner
|
||||
.request::<types::ListToolsResponse>(types::RequestType::ListTools.as_str(), ())
|
||||
.await?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Executes a tool with the given arguments
|
||||
pub async fn run_tool<P: AsRef<str>>(
|
||||
&self,
|
||||
tool: P,
|
||||
arguments: Option<HashMap<String, serde_json::Value>>,
|
||||
) -> Result<types::CallToolResponse> {
|
||||
self.check_capability(ServerCapability::Tools)?;
|
||||
|
||||
let params = types::CallToolParams {
|
||||
name: tool.as_ref().to_string(),
|
||||
arguments,
|
||||
meta: None,
|
||||
};
|
||||
|
||||
let response: types::CallToolResponse = self
|
||||
.inner
|
||||
.request(types::RequestType::CallTool.as_str(), params)
|
||||
.await?;
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
|
||||
impl InitializedContextServerProtocol {
|
||||
pub async fn request<R: serde::de::DeserializeOwned>(
|
||||
&self,
|
||||
method: &str,
|
||||
params: impl serde::Serialize,
|
||||
) -> Result<R> {
|
||||
self.inner.request(method, params).await
|
||||
pub async fn request<T: Request>(&self, params: T::Params) -> Result<T::Response> {
|
||||
self.inner.request(T::METHOD, params).await
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue