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(
|
||||
&self,
|
||||
_request: Option<rmcp::model::PaginatedRequestParams>,
|
||||
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> Result<rmcp::model::ListPromptsResult, rmcp::ErrorData> {
|
||||
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 {
|
||||
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
||||
prompts,
|
||||
meta: #meta,
|
||||
next_cursor: None,
|
||||
ttl_ms: None,
|
||||
cache_scope: None,
|
||||
ttl_ms: supports_cache_hints.then_some(0),
|
||||
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(
|
||||
&self,
|
||||
_request: Option<rmcp::model::PaginatedRequestParams>,
|
||||
_context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
context: rmcp::service::RequestContext<rmcp::RoleServer>,
|
||||
) -> 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{
|
||||
result_type: Some(rmcp::model::ResultType::COMPLETE),
|
||||
tools: #router.list_all(),
|
||||
meta: #result_meta,
|
||||
next_cursor: None,
|
||||
ttl_ms: None,
|
||||
cache_scope: None,
|
||||
ttl_ms: supports_cache_hints.then_some(0),
|
||||
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