fix: emit cache hints from handler macros (#1120)
This commit is contained in:
parent
f8f9607d0b
commit
07bcda2973
3 changed files with 120 additions and 6 deletions
|
|
@ -57,16 +57,20 @@ pub fn prompt_handler(attr: TokenStream, input: TokenStream) -> syn::Result<Toke
|
||||||
async fn list_prompts(
|
async fn list_prompts(
|
||||||
&self,
|
&self,
|
||||||
_request: Option<rmcp::model::PaginatedRequestParams>,
|
_request: Option<rmcp::model::PaginatedRequestParams>,
|
||||||
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||||
) -> Result<rmcp::model::ListPromptsResult, rmcp::ErrorData> {
|
) -> Result<rmcp::model::ListPromptsResult, rmcp::ErrorData> {
|
||||||
let prompts = #router_expr.list_all();
|
let prompts = #router_expr.list_all();
|
||||||
|
let supports_cache_hints = context.protocol_version().is_some_and(|version| {
|
||||||
|
version >= rmcp::model::ProtocolVersion::V_2026_07_28
|
||||||
|
});
|
||||||
Ok(rmcp::model::ListPromptsResult {
|
Ok(rmcp::model::ListPromptsResult {
|
||||||
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
||||||
prompts,
|
prompts,
|
||||||
meta: #meta,
|
meta: #meta,
|
||||||
next_cursor: None,
|
next_cursor: None,
|
||||||
ttl_ms: None,
|
ttl_ms: supports_cache_hints.then_some(0),
|
||||||
cache_scope: None,
|
cache_scope: supports_cache_hints
|
||||||
|
.then_some(rmcp::model::CacheScope::Public),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
|
||||||
|
|
@ -66,15 +66,19 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result<TokenS
|
||||||
async fn list_tools(
|
async fn list_tools(
|
||||||
&self,
|
&self,
|
||||||
_request: Option<rmcp::model::PaginatedRequestParams>,
|
_request: Option<rmcp::model::PaginatedRequestParams>,
|
||||||
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||||
) -> Result<rmcp::model::ListToolsResult, rmcp::ErrorData> {
|
) -> Result<rmcp::model::ListToolsResult, rmcp::ErrorData> {
|
||||||
|
let supports_cache_hints = context.protocol_version().is_some_and(|version| {
|
||||||
|
version >= rmcp::model::ProtocolVersion::V_2026_07_28
|
||||||
|
});
|
||||||
Ok(rmcp::model::ListToolsResult{
|
Ok(rmcp::model::ListToolsResult{
|
||||||
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
||||||
tools: #router.list_all(),
|
tools: #router.list_all(),
|
||||||
meta: #result_meta,
|
meta: #result_meta,
|
||||||
next_cursor: None,
|
next_cursor: None,
|
||||||
ttl_ms: None,
|
ttl_ms: supports_cache_hints.then_some(0),
|
||||||
cache_scope: None,
|
cache_scope: supports_cache_hints
|
||||||
|
.then_some(rmcp::model::CacheScope::Public),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
})?;
|
})?;
|
||||||
|
|
|
||||||
106
crates/rmcp/tests/test_handler_cache_hints.rs
Normal file
106
crates/rmcp/tests/test_handler_cache_hints.rs
Normal file
|
|
@ -0,0 +1,106 @@
|
||||||
|
#![cfg(not(feature = "local"))]
|
||||||
|
#![cfg(feature = "client")]
|
||||||
|
|
||||||
|
use rmcp::{
|
||||||
|
ClientHandler, ServerHandler, ServiceExt,
|
||||||
|
handler::server::router::{prompt::PromptRouter, tool::ToolRouter},
|
||||||
|
model::{CacheScope, ClientInfo, ListPromptsResult, ListToolsResult, ProtocolVersion},
|
||||||
|
prompt_handler, tool_handler,
|
||||||
|
};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct CacheHintServer {
|
||||||
|
tool_router: ToolRouter<Self>,
|
||||||
|
prompt_router: PromptRouter<Self>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CacheHintServer {
|
||||||
|
fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
tool_router: ToolRouter::new(),
|
||||||
|
prompt_router: PromptRouter::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tool_handler(router = self.tool_router)]
|
||||||
|
#[prompt_handler(router = self.prompt_router)]
|
||||||
|
impl ServerHandler for CacheHintServer {}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct VersionedClient {
|
||||||
|
protocol_version: ProtocolVersion,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ClientHandler for VersionedClient {
|
||||||
|
fn get_info(&self) -> ClientInfo {
|
||||||
|
let mut info = ClientInfo::default();
|
||||||
|
info.protocol_version = self.protocol_version.clone();
|
||||||
|
info
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_results(protocol_version: ProtocolVersion) -> (ListToolsResult, ListPromptsResult) {
|
||||||
|
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
CacheHintServer::new()
|
||||||
|
.serve(server_transport)
|
||||||
|
.await?
|
||||||
|
.waiting()
|
||||||
|
.await?;
|
||||||
|
anyhow::Ok(())
|
||||||
|
});
|
||||||
|
|
||||||
|
let client = VersionedClient { protocol_version }
|
||||||
|
.serve(client_transport)
|
||||||
|
.await
|
||||||
|
.expect("client should connect");
|
||||||
|
let tools = client
|
||||||
|
.list_tools(None)
|
||||||
|
.await
|
||||||
|
.expect("tools/list should succeed");
|
||||||
|
let prompts = client
|
||||||
|
.list_prompts(None)
|
||||||
|
.await
|
||||||
|
.expect("prompts/list should succeed");
|
||||||
|
|
||||||
|
client.cancel().await.expect("client should cancel");
|
||||||
|
server_handle.await.expect("server task").expect("server");
|
||||||
|
(tools, prompts)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn handler_macros_should_emit_required_cache_hints_for_2026_07_28() {
|
||||||
|
let (tools, prompts) = list_results(ProtocolVersion::V_2026_07_28).await;
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
(
|
||||||
|
tools.ttl_ms,
|
||||||
|
tools.cache_scope,
|
||||||
|
prompts.ttl_ms,
|
||||||
|
prompts.cache_scope,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
Some(0),
|
||||||
|
Some(CacheScope::Public),
|
||||||
|
Some(0),
|
||||||
|
Some(CacheScope::Public),
|
||||||
|
)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn handler_macros_should_omit_cache_hints_for_legacy_versions() {
|
||||||
|
let (tools, prompts) = list_results(ProtocolVersion::V_2025_11_25).await;
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
(
|
||||||
|
tools.ttl_ms,
|
||||||
|
tools.cache_scope,
|
||||||
|
prompts.ttl_ms,
|
||||||
|
prompts.cache_scope,
|
||||||
|
),
|
||||||
|
(None, None, None, None)
|
||||||
|
);
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue