Skip to content

Commit 1a4a52a

Browse files
DaleSeoDaleSeo
authored
feat: add local feature for !Send tool handler support (#740)
* feat: add local feature for !Send tool handler support * fix: gate streamable HTTP transport on not(local) feature
1 parent 8700e5c commit 1a4a52a

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+362
-222
lines changed

‎crates/rmcp-macros/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,6 @@ serde_json = "1.0"
2222
darling = { version = "0.23" }
2323

2424
[features]
25+
local = []
26+
2527
[dev-dependencies]

‎crates/rmcp-macros/src/tool.rs‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,9 @@ pub struct ToolAttribute {
9595
pub icons: Option<Expr>,
9696
/// Optional metadata for the tool
9797
pub meta: Option<Expr>,
98+
/// When true, the generated future will not require `Send`. Useful for `!Send` handlers
99+
/// (e.g. single-threaded database connections). Also enabled globally by the `local` crate feature.
100+
pub local: bool,
98101
}
99102

100103
#[derive(FromMeta, Debug, Default)]
@@ -333,7 +336,9 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
333336
if fn_item.sig.asyncness.is_some() {
334337
// 1. remove asyncness from sig
335338
// 2. make return type: `std::pin::Pin<Box<dyn std::future::Future<Output = #ReturnType> + Send + '_>>`
339+
// (omit `+ Send` when the `local` crate feature is active or `#[tool(local)]` is used)
336340
// 3. make body: { Box::pin(async move { #body }) }
341+
let omit_send = cfg!(feature = "local") || attribute.local;
337342
let new_output = syn::parse2::<ReturnType>({
338343
let mut lt = quote! { 'static };
339344
if let Some(receiver) = fn_item.sig.receiver() {
@@ -347,10 +352,18 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
347352
}
348353
match &fn_item.sig.output {
349354
syn::ReturnType::Default => {
350-
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = ()> + Send + #lt>> }
355+
if omit_send {
356+
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = ()> + #lt>> }
357+
} else {
358+
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = ()> + Send + #lt>> }
359+
}
351360
}
352361
syn::ReturnType::Type(_, ty) => {
353-
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = #ty> + Send + #lt>> }
362+
if omit_send {
363+
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = #ty> + #lt>> }
364+
} else {
365+
quote! { -> ::std::pin::Pin<Box<dyn ::std::future::Future<Output = #ty> + Send + #lt>> }
366+
}
354367
}
355368
}
356369
})?;

‎crates/rmcp/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ chrono = { version = "0.4.38", default-features = false, features = [
7777

7878
[features]
7979
default = ["base64", "macros", "server"]
80+
local = ["rmcp-macros?/local"]
8081
client = ["dep:tokio-stream"]
8182
server = ["transport-async-rw", "dep:schemars", "dep:pastey"]
8283
macros = ["dep:rmcp-macros", "dep:pastey"]

‎crates/rmcp/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ The transport layer is pluggable. Two built-in pairs cover the most common cases
5252
| | Client | Server |
5353
|:-:|:-:|:-:|
5454
| **stdio** | [`TokioChildProcess`](crate::transport::TokioChildProcess) | [`stdio`](crate::transport::stdio) |
55-
| **Streamable HTTP** | [`StreamableHttpClientTransport`](crate::transport::StreamableHttpClientTransport) | [`StreamableHttpService`](crate::transport::StreamableHttpService) |
55+
| **Streamable HTTP** | [`StreamableHttpClientTransport`](crate::transport::StreamableHttpClientTransport) | `StreamableHttpService` |
5656

5757
Any type that implements the [`Transport`](crate::transport::Transport) trait can be used. The [`IntoTransport`](crate::transport::IntoTransport) helper trait provides automatic conversions from:
5858

‎crates/rmcp/src/handler/client.rs‎

Lines changed: 31 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ use std::sync::Arc;
44
use crate::{
55
error::ErrorData as McpError,
66
model::*,
7-
service::{NotificationContext, RequestContext, RoleClient, Service, ServiceRole},
7+
service::{
8+
MaybeSendFuture, NotificationContext, RequestContext, RoleClient, Service, ServiceRole,
9+
},
810
};
911

1012
impl<H: ClientHandler> Service<RoleClient> for H {
@@ -83,15 +85,15 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
8385
fn ping(
8486
&self,
8587
context: RequestContext<RoleClient>,
86-
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
88+
) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_ {
8789
std::future::ready(Ok(()))
8890
}
8991

9092
fn create_message(
9193
&self,
9294
params: CreateMessageRequestParams,
9395
context: RequestContext<RoleClient>,
94-
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + Send + '_ {
96+
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + MaybeSendFuture + '_ {
9597
std::future::ready(Err(
9698
McpError::method_not_found::<CreateMessageRequestMethod>(),
9799
))
@@ -100,7 +102,7 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
100102
fn list_roots(
101103
&self,
102104
context: RequestContext<RoleClient>,
103-
) -> impl Future<Output = Result<ListRootsResult, McpError>> + Send + '_ {
105+
) -> impl Future<Output = Result<ListRootsResult, McpError>> + MaybeSendFuture + '_ {
104106
std::future::ready(Ok(ListRootsResult::default()))
105107
}
106108

@@ -162,7 +164,8 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
162164
&self,
163165
request: CreateElicitationRequestParams,
164166
context: RequestContext<RoleClient>,
165-
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + Send + '_ {
167+
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + MaybeSendFuture + '_
168+
{
166169
// Default implementation declines all requests - real clients should override this
167170
let _ = (request, context);
168171
std::future::ready(Ok(CreateElicitationResult {
@@ -175,7 +178,7 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
175178
&self,
176179
request: CustomRequest,
177180
context: RequestContext<RoleClient>,
178-
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
181+
) -> impl Future<Output = Result<CustomResult, McpError>> + MaybeSendFuture + '_ {
179182
let CustomRequest { method, .. } = request;
180183
let _ = context;
181184
std::future::ready(Err(McpError::new(
@@ -189,61 +192,61 @@ pub trait ClientHandler: Sized + Send + Sync + 'static {
189192
&self,
190193
params: CancelledNotificationParam,
191194
context: NotificationContext<RoleClient>,
192-
) -> impl Future<Output = ()> + Send + '_ {
195+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
193196
std::future::ready(())
194197
}
195198
fn on_progress(
196199
&self,
197200
params: ProgressNotificationParam,
198201
context: NotificationContext<RoleClient>,
199-
) -> impl Future<Output = ()> + Send + '_ {
202+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
200203
std::future::ready(())
201204
}
202205
fn on_logging_message(
203206
&self,
204207
params: LoggingMessageNotificationParam,
205208
context: NotificationContext<RoleClient>,
206-
) -> impl Future<Output = ()> + Send + '_ {
209+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
207210
std::future::ready(())
208211
}
209212
fn on_resource_updated(
210213
&self,
211214
params: ResourceUpdatedNotificationParam,
212215
context: NotificationContext<RoleClient>,
213-
) -> impl Future<Output = ()> + Send + '_ {
216+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
214217
std::future::ready(())
215218
}
216219
fn on_resource_list_changed(
217220
&self,
218221
context: NotificationContext<RoleClient>,
219-
) -> impl Future<Output = ()> + Send + '_ {
222+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
220223
std::future::ready(())
221224
}
222225
fn on_tool_list_changed(
223226
&self,
224227
context: NotificationContext<RoleClient>,
225-
) -> impl Future<Output = ()> + Send + '_ {
228+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
226229
std::future::ready(())
227230
}
228231
fn on_prompt_list_changed(
229232
&self,
230233
context: NotificationContext<RoleClient>,
231-
) -> impl Future<Output = ()> + Send + '_ {
234+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
232235
std::future::ready(())
233236
}
234237

235238
fn on_url_elicitation_notification_complete(
236239
&self,
237240
params: ElicitationResponseNotificationParam,
238241
context: NotificationContext<RoleClient>,
239-
) -> impl Future<Output = ()> + Send + '_ {
242+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
240243
std::future::ready(())
241244
}
242245
fn on_custom_notification(
243246
&self,
244247
notification: CustomNotification,
245248
context: NotificationContext<RoleClient>,
246-
) -> impl Future<Output = ()> + Send + '_ {
249+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
247250
let _ = (notification, context);
248251
std::future::ready(())
249252
}
@@ -269,99 +272,99 @@ macro_rules! impl_client_handler_for_wrapper {
269272
fn ping(
270273
&self,
271274
context: RequestContext<RoleClient>,
272-
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
275+
) -> impl Future<Output = Result<(), McpError>> + MaybeSendFuture + '_ {
273276
(**self).ping(context)
274277
}
275278

276279
fn create_message(
277280
&self,
278281
params: CreateMessageRequestParams,
279282
context: RequestContext<RoleClient>,
280-
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + Send + '_ {
283+
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + MaybeSendFuture + '_ {
281284
(**self).create_message(params, context)
282285
}
283286

284287
fn list_roots(
285288
&self,
286289
context: RequestContext<RoleClient>,
287-
) -> impl Future<Output = Result<ListRootsResult, McpError>> + Send + '_ {
290+
) -> impl Future<Output = Result<ListRootsResult, McpError>> + MaybeSendFuture + '_ {
288291
(**self).list_roots(context)
289292
}
290293

291294
fn create_elicitation(
292295
&self,
293296
request: CreateElicitationRequestParams,
294297
context: RequestContext<RoleClient>,
295-
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + Send + '_ {
298+
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + MaybeSendFuture + '_ {
296299
(**self).create_elicitation(request, context)
297300
}
298301

299302
fn on_custom_request(
300303
&self,
301304
request: CustomRequest,
302305
context: RequestContext<RoleClient>,
303-
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
306+
) -> impl Future<Output = Result<CustomResult, McpError>> + MaybeSendFuture + '_ {
304307
(**self).on_custom_request(request, context)
305308
}
306309

307310
fn on_cancelled(
308311
&self,
309312
params: CancelledNotificationParam,
310313
context: NotificationContext<RoleClient>,
311-
) -> impl Future<Output = ()> + Send + '_ {
314+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
312315
(**self).on_cancelled(params, context)
313316
}
314317

315318
fn on_progress(
316319
&self,
317320
params: ProgressNotificationParam,
318321
context: NotificationContext<RoleClient>,
319-
) -> impl Future<Output = ()> + Send + '_ {
322+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
320323
(**self).on_progress(params, context)
321324
}
322325

323326
fn on_logging_message(
324327
&self,
325328
params: LoggingMessageNotificationParam,
326329
context: NotificationContext<RoleClient>,
327-
) -> impl Future<Output = ()> + Send + '_ {
330+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
328331
(**self).on_logging_message(params, context)
329332
}
330333

331334
fn on_resource_updated(
332335
&self,
333336
params: ResourceUpdatedNotificationParam,
334337
context: NotificationContext<RoleClient>,
335-
) -> impl Future<Output = ()> + Send + '_ {
338+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
336339
(**self).on_resource_updated(params, context)
337340
}
338341

339342
fn on_resource_list_changed(
340343
&self,
341344
context: NotificationContext<RoleClient>,
342-
) -> impl Future<Output = ()> + Send + '_ {
345+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
343346
(**self).on_resource_list_changed(context)
344347
}
345348

346349
fn on_tool_list_changed(
347350
&self,
348351
context: NotificationContext<RoleClient>,
349-
) -> impl Future<Output = ()> + Send + '_ {
352+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
350353
(**self).on_tool_list_changed(context)
351354
}
352355

353356
fn on_prompt_list_changed(
354357
&self,
355358
context: NotificationContext<RoleClient>,
356-
) -> impl Future<Output = ()> + Send + '_ {
359+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
357360
(**self).on_prompt_list_changed(context)
358361
}
359362

360363
fn on_custom_notification(
361364
&self,
362365
notification: CustomNotification,
363366
context: NotificationContext<RoleClient>,
364-
) -> impl Future<Output = ()> + Send + '_ {
367+
) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
365368
(**self).on_custom_notification(notification, context)
366369
}
367370

0 commit comments

Comments
 (0)