tools_kit/sources/mcp/server/
mod.rs

1mod serving;
2mod sessions;
3
4use std::{sync::Arc, time::Duration};
5
6use axum::{
7    Router,
8    extract::Request,
9    http::{StatusCode, header},
10    middleware::{self, Next},
11    response::{IntoResponse, Response},
12};
13use rmcp::{
14    ErrorData as McpError, RoleServer, ServerHandler,
15    model::{
16        CallToolRequestParams, CallToolResult, ClientResult, CreateElicitationRequestParams,
17        ElicitationAction, ElicitationSchema, ListToolsResult, PaginatedRequestParams,
18        ServerCapabilities, ServerInfo, ServerRequest,
19    },
20    service::{ElicitationMode, PeerRequestOptions, RequestContext},
21    transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService},
22};
23use tokio::{net::TcpListener, time};
24use url::{Host, Url};
25
26use crate::{
27    ToolError, ToolHandlerBackend, ToolsetBackend,
28    declaration::ToolInfo,
29    execution::{ToolHandlerInput, ToolHandlerOutput, ToolInvocation},
30    interaction::{ToolInteractionKind, ToolInteractionRequest, ToolInteractionResponse},
31    mcp::{McpServerConfig, McpServerInteractionPolicy},
32    name::ToolName,
33    sources::mcp::catalog::McpProtocol,
34};
35
36use self::{
37    serving::Serving,
38    sessions::{Protocol, Sessions},
39};
40
41const LIST_TOOLS_PAGE_SIZE: usize = 64;
42const INTERACTION_TIMEOUT: Duration = Duration::from_secs(30);
43const INTERACTION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(1);
44
45/// Independently owned MCP server serving any concrete [`ToolsetBackend`].
46///
47/// The server retains a shared reference to the caller's toolset. Server shutdown
48/// never shuts down that toolset; the caller owns its cleanup separately.
49/// This example binds a real loopback endpoint and requires an available port.
50/// Keep the Tokio runtime driven until both server and toolset cleanup settle.
51/// Dropping the server or its shutdown observation signals its retained cleanup owner;
52/// only explicit toolset shutdown closes the separate executable source.
53///
54/// # Examples
55///
56/// ```no_run
57/// use std::sync::Arc;
58/// use tools_kit::{Toolset, ToolsetBackend, ToolHandlerBackend, ToolError};
59/// use tools_kit::declaration::ToolName;
60/// use tools_kit::execution::{ToolHandlerInput, ToolHandlerOutput, ToolInvocation};
61/// use tools_kit::local::ToolRegistration;
62/// use tools_kit::mcp::{McpServer, McpServerConfig, McpClientConfig, McpConnectionConfig};
63/// // MCP requires object-shaped arguments; local typed tools may use scalars.
64/// #[derive(serde::Deserialize, schemars::JsonSchema)]
65/// #[serde(deny_unknown_fields)]
66/// struct DoubleArgs { value: i64 }
67/// #[derive(serde::Serialize, schemars::JsonSchema)]
68/// struct DoubleResult { value: i64 }
69/// # #[tokio::main(flavor = "current_thread")]
70/// # async fn main() -> Result<(), ToolError> {
71/// let source = Arc::new(Toolset::local(vec![ToolRegistration::typed(
72///     "double", |input: DoubleArgs| async move {
73///         let value = input.value.checked_mul(2).ok_or(ToolError::InvalidArguments {
74///             tool: ToolName::new("double")?, message: "doubling would overflow i64".into(), issues: Vec::new(),
75///         })?;
76///         Ok(DoubleResult { value })
77///     },
78/// )?])?);
79/// let config = McpServerConfig::new("127.0.0.1:8000".parse().unwrap(), "/mcp", "page")?;
80/// let server = McpServer::start(Arc::clone(&source), config).await;
81/// let observed = async { match server {
82///     Ok(server) => {
83///         let client_work = async {
84///             let client = Toolset::mcp(McpClientConfig::new(McpConnectionConfig::StreamableHttp {
85///                 endpoint: "http://127.0.0.1:8000/mcp".parse().unwrap(), headers: Default::default(),
86///             })).await?;
87///             let invoked = async {
88///                 let mut handler = client.tool_handler(ToolName::new("double")?).await?;
89///                 let output = handler.step(ToolHandlerInput::Invoke(ToolInvocation {
90///                     call_id: Some("client-call".into()), arguments: serde_json::json!({"value":21}),
91///                 })).await?;
92///                 let ToolHandlerOutput::Completed(result) = output else {
93///                     return Err(ToolError::Protocol { message: "unexpected interaction".into() });
94///                 };
95///                 assert!(!result.is_error);
96///                 assert_eq!(result.structured_content, Some(serde_json::json!({"value":42})));
97///                 Ok::<(), ToolError>(())
98///             }.await;
99///             let client_closed = client.shutdown().await;
100///             invoked?;
101///             client_closed
102///         }.await;
103///         let serving_closed = server.shutdown().await; // Consumes and joins serving.
104///         client_work?;
105///         serving_closed?;
106///         assert_eq!(source.list_tools().await?.len(), 1); // Source remains usable.
107///         Ok(())
108///     }
109///     Err(error) => Err(error),
110/// }
111/// }.await;
112/// let source_closed = source.shutdown().await;
113/// observed?;
114/// source_closed?;
115/// # Ok(())
116/// # }
117/// ```
118///
119pub struct McpServer {
120    serving: Arc<Serving>,
121}
122
123impl McpServer {
124    /// Binds and starts the configured loopback endpoint eagerly.
125    ///
126    /// A failed start leaves the caller's toolset usable. Neither the toolset nor
127    /// its associated handler needs to implement `Clone` or use supplied handles.
128    ///
129    /// This example binds a real loopback endpoint and requires an available port.
130    /// Keep the Tokio runtime driven through both independently owned cleanup operations.
131    ///
132    /// ```no_run
133    /// # #[tokio::main(flavor = "current_thread")]
134    /// # async fn main() -> Result<(), tools_kit::ToolError> {
135    /// use std::sync::Arc;
136    /// use tools_kit::{Toolset, ToolsetBackend, mcp::{McpServer, McpServerConfig}};
137    /// let tools = Arc::new(Toolset::local(Vec::new())?);
138    /// let config = McpServerConfig::new("127.0.0.1:8000".parse().unwrap(), "/mcp", "page")?;
139    /// let result = McpServer::start(Arc::clone(&tools), config).await;
140    /// let stopped = match result {
141    ///     Ok(server) => server.shutdown().await,
142    ///     Err(error) => Err(error),
143    /// };
144    /// let source_closed = tools.shutdown().await;
145    /// stopped?;
146    /// source_closed?;
147    /// # Ok(())
148    /// # }
149    /// ```
150    ///
151    /// # Errors
152    ///
153    /// Returns [`ToolError::InvalidConfiguration`] for invalid server settings or [`ToolError::Transport`] when the endpoint cannot bind. Failed startup leaves source shutdown with the caller.
154    ///
155    pub async fn start<T: ToolsetBackend>(
156        toolset: Arc<T>,
157        config: McpServerConfig,
158    ) -> Result<Self, ToolError> {
159        config.validate()?;
160        let listener =
161            TcpListener::bind(config.bind())
162                .await
163                .map_err(|error| ToolError::Transport {
164                    tool: None,
165                    message: format!("failed to bind MCP server endpoint: {error}"),
166                })?;
167        let serving = Serving::new(config.shutdown_timeout());
168        let protocol = Arc::new(Protocol::new(Arc::clone(&serving)));
169        let handler = McpServeHandler {
170            serving: Arc::clone(&serving),
171            toolset,
172            interaction_policy: config.interaction_policy(),
173            cursor_prefix: config.cursor_prefix().to_owned(),
174            _protocol: Arc::clone(&protocol),
175        };
176        let transport_config = StreamableHttpServerConfig::default();
177        let stream_cancellation = transport_config.cancellation_token.clone();
178        let service: StreamableHttpService<McpServeHandler<T>, Sessions> =
179            StreamableHttpService::new(
180                move || {
181                    if handler.serving.closing() {
182                        Err(std::io::Error::other("MCP server is closing"))
183                    } else {
184                        Ok(handler.clone())
185                    }
186                },
187                Arc::new(Sessions::new(Arc::clone(&serving), protocol)),
188                transport_config,
189            );
190        let router = Router::new()
191            .nest_service(config.endpoint_path(), service)
192            .layer(middleware::from_fn(McpOriginPolicy::enforce));
193        let stopping = Arc::clone(&serving);
194        serving.start(async move {
195            axum::serve(listener, router)
196                .with_graceful_shutdown(async move {
197                    stopping.stopped().await;
198                    stream_cancellation.cancel();
199                })
200                .await
201        });
202        Ok(Self { serving })
203    }
204
205    /// Seals admission and observes the retained server owner through cleanup.
206    ///
207    /// A timeout reports unsettled cleanup; it does not abort the cleanup owner.
208    /// This example binds a real loopback endpoint and requires an available port.
209    /// Keep the Tokio runtime driven until both server and toolset cleanup settle.
210    /// Dropping the server or its shutdown observation signals its retained cleanup owner;
211    /// only explicit toolset shutdown closes the separate executable source.
212    ///
213    /// # Errors
214    ///
215    /// Returns the retained server cleanup failure, including timeout or handler-cleanup failure. A timeout leaves cleanup owned rather than aborting it.
216    ///
217    /// # Examples
218    ///
219    /// ```no_run
220    ///
221    /// # #[tokio::main(flavor = "current_thread")]
222    /// # async fn main() -> Result<(), tools_kit::ToolError> {
223    /// use std::sync::Arc;
224    /// use tools_kit::{Toolset, ToolsetBackend, mcp::{McpServer, McpServerConfig}};
225    /// let tools = Arc::new(Toolset::local(Vec::new())?);
226    /// let config = McpServerConfig::new("127.0.0.1:8000".parse().unwrap(), "/mcp", "page")?;
227    /// let server = McpServer::start(Arc::clone(&tools), config).await?;
228    /// let stopped = server.shutdown().await;
229    /// let source_closed = tools.shutdown().await;
230    /// stopped?;
231    /// source_closed?;
232    ///
233    /// # Ok(())
234    /// # }
235    /// ```
236    ///
237    pub async fn shutdown(self) -> Result<(), ToolError> {
238        self.serving.shutdown().await
239    }
240}
241
242impl Drop for McpServer {
243    fn drop(&mut self) {
244        self.serving.release_observer();
245    }
246}
247
248impl std::fmt::Debug for McpServer {
249    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
250        formatter.debug_struct("McpServer").finish_non_exhaustive()
251    }
252}
253
254struct McpServeHandler<T: ToolsetBackend> {
255    serving: Arc<Serving>,
256    toolset: Arc<T>,
257    interaction_policy: McpServerInteractionPolicy,
258    cursor_prefix: String,
259    // Final field: retire protocol ownership after caller resources.
260    _protocol: Arc<Protocol>,
261}
262
263impl<T: ToolsetBackend> Clone for McpServeHandler<T> {
264    fn clone(&self) -> Self {
265        Self {
266            serving: Arc::clone(&self.serving),
267            toolset: Arc::clone(&self.toolset),
268            interaction_policy: self.interaction_policy,
269            cursor_prefix: self.cursor_prefix.clone(),
270            _protocol: Arc::clone(&self._protocol),
271        }
272    }
273}
274
275impl<T: ToolsetBackend> ServerHandler for McpServeHandler<T> {
276    fn get_info(&self) -> ServerInfo {
277        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
278    }
279
280    async fn list_tools(
281        &self,
282        request: Option<PaginatedRequestParams>,
283        context: RequestContext<RoleServer>,
284    ) -> Result<ListToolsResult, McpError> {
285        let handler = self.clone();
286        let owner = Arc::clone(&self.serving);
287        let cancellation = context.ct.clone();
288        let observation = self.serving.admit(async move {
289            let operation = handler.list_tools_owned(request, context);
290            tokio::pin!(operation);
291            tokio::select! {
292                biased;
293                _ = owner.stopped() => {
294                    cancellation.cancel();
295                    operation.await
296                }
297                result = &mut operation => result,
298            }
299        })?;
300        observation
301            .await
302            .map_err(|_| McpError::internal_error("MCP server inventory task terminated", None))?
303    }
304
305    async fn call_tool(
306        &self,
307        request: CallToolRequestParams,
308        context: RequestContext<RoleServer>,
309    ) -> Result<CallToolResult, McpError> {
310        let handler = self.clone();
311        let owner = Arc::clone(&self.serving);
312        let cancellation = context.ct.clone();
313        let observation = self.serving.admit(async move {
314            let operation = handler.call_tool_owned(request, context);
315            tokio::pin!(operation);
316            tokio::select! {
317                biased;
318                _ = owner.stopped() => {
319                    cancellation.cancel();
320                    operation.await
321                }
322                result = &mut operation => result,
323            }
324        })?;
325        observation
326            .await
327            .map_err(|_| McpError::internal_error("MCP server invocation task terminated", None))?
328    }
329}
330
331impl<T: ToolsetBackend> McpServeHandler<T> {
332    async fn list_tools_owned(
333        &self,
334        request: Option<PaginatedRequestParams>,
335        context: RequestContext<RoleServer>,
336    ) -> Result<ListToolsResult, McpError> {
337        if context.ct.is_cancelled() {
338            return Err(McpError::internal_error(
339                "MCP tools/list request was cancelled",
340                None,
341            ));
342        }
343        let tools = self
344            .toolset
345            .list_tools()
346            .await
347            .map_err(|error| McpError::internal_error(error.to_string(), None))?;
348        let start = self.list_tools_offset(request, &context, tools.len())?;
349        let page = tools
350            .iter()
351            .skip(start)
352            .take(LIST_TOOLS_PAGE_SIZE)
353            .map(McpProtocol::tool_info_to_mcp)
354            .collect::<Result<Vec<_>, _>>()
355            .map_err(|error| McpError::internal_error(error.to_string(), None))?;
356        let next_offset = start + page.len();
357        let mut result = ListToolsResult::with_all_items(page);
358        result.next_cursor =
359            (next_offset < tools.len()).then(|| format!("{}{next_offset}", self.cursor_prefix));
360        Ok(result)
361    }
362
363    async fn call_tool_owned(
364        &self,
365        request: CallToolRequestParams,
366        context: RequestContext<RoleServer>,
367    ) -> Result<CallToolResult, McpError> {
368        let name = ToolName::new(request.name.to_string())
369            .map_err(|error| McpError::invalid_params(error.to_string(), None))?;
370        self.check_cancellation(&name, &context)
371            .map_err(McpProtocol::tool_error_to_mcp_error)?;
372        let info = self
373            .declaration(&name)
374            .await
375            .map_err(McpProtocol::tool_error_to_mcp_error)?;
376        self.check_cancellation(&name, &context)
377            .map_err(McpProtocol::tool_error_to_mcp_error)?;
378        let arguments = serde_json::Value::Object(request.arguments.unwrap_or_default());
379        let issues = info
380            .input_schema
381            .validate(&arguments)
382            .map_err(McpProtocol::tool_error_to_mcp_error)?;
383        if !issues.is_empty() {
384            return Err(McpProtocol::tool_error_to_mcp_error(
385                ToolError::InvalidArguments {
386                    tool: name,
387                    message: "invocation failed the declared input schema".to_owned(),
388                    issues,
389                },
390            ));
391        }
392        let mut handler = self
393            .toolset
394            .tool_handler(name.clone())
395            .await
396            .map_err(McpProtocol::tool_error_to_mcp_error)?;
397        // Construction may acquire caller-owned resources. Retain it even when
398        // cancelled, then cancel the returned handler before its first Invoke.
399        let mut output = self
400            .step(
401                &mut handler,
402                ToolHandlerInput::Invoke(ToolInvocation {
403                    call_id: None,
404                    arguments,
405                }),
406                &name,
407                &context,
408            )
409            .await
410            .map_err(McpProtocol::tool_error_to_mcp_error)?;
411        loop {
412            match output {
413                ToolHandlerOutput::Completed(result) => {
414                    info.validate_output(&result)
415                        .map_err(McpProtocol::tool_error_to_mcp_error)?;
416                    return McpProtocol::tool_result_to_mcp(result)
417                        .map_err(|error| McpError::internal_error(error.to_string(), None));
418                }
419                ToolHandlerOutput::Interaction(request) => {
420                    let response = self
421                        .interaction_policy
422                        .resolve(request, &context, &self.serving)
423                        .await;
424                    let response = match response {
425                        Ok(response) => response,
426                        Err(primary) => {
427                            return Err(McpProtocol::tool_error_to_mcp_error(
428                                self.cancel(&mut handler, &name, primary).await,
429                            ));
430                        }
431                    };
432                    output = self
433                        .step(
434                            &mut handler,
435                            ToolHandlerInput::Interaction(response),
436                            &name,
437                            &context,
438                        )
439                        .await
440                        .map_err(McpProtocol::tool_error_to_mcp_error)?;
441                }
442            }
443        }
444    }
445
446    fn check_cancellation(
447        &self,
448        name: &ToolName,
449        context: &RequestContext<RoleServer>,
450    ) -> Result<(), ToolError> {
451        if context.ct.is_cancelled() {
452            Err(ToolError::Cancelled {
453                tool: name.clone(),
454                reason: Some("MCP invocation was cancelled".to_owned()),
455            })
456        } else {
457            Ok(())
458        }
459    }
460
461    async fn step(
462        &self,
463        handler: &mut T::Handler,
464        input: ToolHandlerInput,
465        name: &ToolName,
466        context: &RequestContext<RoleServer>,
467    ) -> Result<ToolHandlerOutput, ToolError> {
468        if let Err(primary) = self.check_cancellation(name, context) {
469            return Err(self.cancel(handler, name, primary).await);
470        }
471        let output = {
472            let step = handler.step(input);
473            tokio::pin!(step);
474            tokio::select! {
475                biased;
476                _ = context.ct.cancelled() => None,
477                output = &mut step => Some(output),
478            }
479        };
480        // End the pending step's mutable borrow before issuing Cancel on the
481        // same unique handler. Await cleanup rather than cancelling its waiter.
482        match output {
483            Some(output) => output,
484            None => Err(self
485                .cancel(
486                    handler,
487                    name,
488                    ToolError::Cancelled {
489                        tool: name.clone(),
490                        reason: Some("MCP invocation was cancelled".to_owned()),
491                    },
492                )
493                .await),
494        }
495    }
496
497    async fn cancel(
498        &self,
499        handler: &mut T::Handler,
500        name: &ToolName,
501        primary: ToolError,
502    ) -> ToolError {
503        match handler
504            .step(ToolHandlerInput::Cancel {
505                reason: Some(primary.to_string()),
506            })
507            .await
508        {
509            Ok(ToolHandlerOutput::Completed(_)) | Err(ToolError::Cancelled { .. }) => primary,
510            Ok(ToolHandlerOutput::Interaction(_)) => {
511                let error = ToolError::Transport {
512                    tool: Some(name.clone()),
513                    message: format!("{primary}; handler cancellation left an interaction pending"),
514                };
515                self.serving.record(error.clone());
516                error
517            }
518            Err(cleanup) => {
519                let error = ToolError::Transport {
520                    tool: Some(name.clone()),
521                    message: format!("{primary}; handler cancellation also failed: {cleanup}"),
522                };
523                self.serving.record(error.clone());
524                error
525            }
526        }
527    }
528
529    async fn declaration(&self, name: &ToolName) -> Result<ToolInfo, ToolError> {
530        let mut matching = self
531            .toolset
532            .list_tools()
533            .await?
534            .into_iter()
535            .filter(|info| &info.name == name);
536        let info = matching
537            .next()
538            .ok_or_else(|| ToolError::UnknownTool { name: name.clone() })?;
539        if matching.next().is_some() {
540            return Err(ToolError::DuplicateName { name: name.clone() });
541        }
542        Ok(info)
543    }
544
545    fn list_tools_offset(
546        &self,
547        request: Option<PaginatedRequestParams>,
548        context: &RequestContext<RoleServer>,
549        tool_count: usize,
550    ) -> Result<usize, McpError> {
551        if context.ct.is_cancelled() {
552            return Err(McpError::internal_error(
553                "MCP tools/list request was cancelled".to_owned(),
554                None,
555            ));
556        }
557        let Some(cursor) = request.and_then(|params| params.cursor) else {
558            return Ok(0);
559        };
560        let Some(raw_offset) = cursor.strip_prefix(&self.cursor_prefix) else {
561            return Err(McpError::invalid_params(
562                format!("invalid MCP tools/list cursor `{cursor}`"),
563                None,
564            ));
565        };
566        let offset = raw_offset.parse::<usize>().map_err(|_| {
567            McpError::invalid_params(format!("invalid MCP tools/list cursor `{cursor}`"), None)
568        })?;
569        if offset == 0 || offset >= tool_count || offset % LIST_TOOLS_PAGE_SIZE != 0 {
570            return Err(McpError::invalid_params(
571                format!("invalid MCP tools/list cursor `{cursor}`"),
572                None,
573            ));
574        }
575        Ok(offset)
576    }
577}
578
579impl McpServerInteractionPolicy {
580    async fn resolve(
581        self,
582        request: ToolInteractionRequest,
583        context: &RequestContext<RoleServer>,
584        serving: &Serving,
585    ) -> Result<ToolInteractionResponse, ToolError> {
586        match (self, &request.kind) {
587            (Self::RejectWithToolError, _) => Err(ToolError::Unsupported {
588                operation: "MCP server tool interaction".to_owned(),
589            }),
590            (Self::ElicitFormIfClientSupports, ToolInteractionKind::Form { schema, .. })
591                if context
592                    .peer
593                    .supported_elicitation_modes()
594                    .contains(&ElicitationMode::Form) =>
595            {
596                let schema = self.elicitation_schema(schema.clone())?;
597                let deadline = time::Instant::now() + INTERACTION_TIMEOUT;
598                let admission = context.peer.send_cancellable_request(
599                    ServerRequest::CreateElicitationRequest(rmcp::model::Request::new(
600                        CreateElicitationRequestParams::FormElicitationParams {
601                            meta: None,
602                            message: request.prompt.clone(),
603                            requested_schema: schema,
604                        },
605                    )),
606                    PeerRequestOptions::no_options(),
607                );
608                // rmcp returns its handle synchronously after queue admission.
609                // Cancelling this wait cannot discard an already admitted handle.
610                let mut elicitation = tokio::select! {
611                    biased;
612                    _ = context.ct.cancelled() => return Err(ToolError::Cancelled {
613                        tool: request.tool_name.clone(),
614                        reason: Some("MCP invocation was cancelled before elicitation admission".to_owned()),
615                    }),
616                    _ = time::sleep_until(deadline) => return Err(ToolError::Timeout {
617                        tool: request.tool_name.clone(),
618                        timeout_secs: INTERACTION_TIMEOUT.as_secs(),
619                    }),
620                    result = admission => result,
621                }
622                .map_err(|error| ToolError::Transport {
623                        tool: Some(request.tool_name.clone()),
624                        message: format!("MCP client elicitation failed: {error}"),
625                    })?;
626                // Keep the native handle until observation or explicit cancellation
627                // completes. Dropping rmcp's response future sends no cancellation.
628                let observed = tokio::select! {
629                    biased;
630                    _ = context.ct.cancelled() => Err(ToolError::Cancelled {
631                        tool: request.tool_name.clone(),
632                        reason: Some("MCP invocation was cancelled during elicitation".to_owned()),
633                    }),
634                    _ = time::sleep_until(deadline) => Err(ToolError::Timeout {
635                        tool: request.tool_name.clone(),
636                        timeout_secs: INTERACTION_TIMEOUT.as_secs(),
637                    }),
638                    result = &mut elicitation.rx => Ok(result),
639                };
640                let response = match observed {
641                    Ok(response) => response,
642                    Err(primary) => {
643                        return Err(match time::timeout(
644                            INTERACTION_CLEANUP_TIMEOUT,
645                            elicitation.cancel(Some(primary.to_string())),
646                        )
647                        .await
648                        {
649                            Ok(Ok(())) => primary,
650                            Ok(Err(cleanup)) => {
651                                let error = ToolError::Transport {
652                                    tool: Some(request.tool_name.clone()),
653                                    message: format!(
654                                        "{primary}; elicitation cancellation also failed: {cleanup}"
655                                    ),
656                                };
657                                serving.record(error.clone());
658                                error
659                            }
660                            Err(_) => {
661                                let error = ToolError::Transport {
662                                    tool: Some(request.tool_name.clone()),
663                                    message: format!(
664                                        "{primary}; elicitation cancellation acknowledgement timed out; client request settlement is unconfirmed"
665                                    ),
666                                };
667                                serving.record(error.clone());
668                                error
669                            }
670                        });
671                    }
672                }
673                .map_err(|error| ToolError::Transport {
674                    tool: Some(request.tool_name.clone()),
675                    message: format!("MCP client elicitation response was lost: {error}"),
676                })?
677                .map_err(|error| ToolError::Transport {
678                    tool: Some(request.tool_name.clone()),
679                    message: format!("MCP client elicitation failed: {error}"),
680                })?;
681                let ClientResult::CreateElicitationResult(result) = response else {
682                    return Err(ToolError::Protocol {
683                        message: "MCP client returned an unexpected elicitation response"
684                            .to_owned(),
685                    });
686                };
687                let id = request.id.clone();
688                let response = match result.action {
689                    ElicitationAction::Accept => ToolInteractionResponse::submit(
690                        id,
691                        result.content.unwrap_or(serde_json::Value::Null),
692                    ),
693                    ElicitationAction::Decline => ToolInteractionResponse::decline(id, None),
694                    ElicitationAction::Cancel => ToolInteractionResponse::cancel(id, None),
695                };
696                // A rejected response can leave the unique handler pending.
697                // Validate before advancing it so the owning call cancels and
698                // awaits that handler through its existing resolution-error path.
699                request.validate_response(&response)?;
700                Ok(response)
701            }
702            (Self::ElicitFormIfClientSupports, ToolInteractionKind::Form { .. }) => {
703                Err(ToolError::Unsupported {
704                    operation: "MCP client form elicitation".to_owned(),
705                })
706            }
707            (Self::ElicitFormIfClientSupports, ToolInteractionKind::Text { .. }) => {
708                Err(ToolError::Unsupported {
709                    operation: "MCP server text interaction".to_owned(),
710                })
711            }
712        }
713    }
714
715    fn elicitation_schema(self, schema: serde_json::Value) -> Result<ElicitationSchema, ToolError> {
716        let serde_json::Value::Object(object) = schema else {
717            return Err(ToolError::Unsupported {
718                operation: "non-object MCP form schema".to_owned(),
719            });
720        };
721        ElicitationSchema::from_json_schema(object).map_err(|error| ToolError::Unsupported {
722            operation: format!("MCP form schema: {error}"),
723        })
724    }
725}
726
727struct McpOriginPolicy;
728
729impl McpOriginPolicy {
730    async fn enforce(request: Request, next: Next) -> Response {
731        if let Some(origin) = request.headers().get(header::ORIGIN)
732            && !Self::is_loopback(origin)
733        {
734            return StatusCode::FORBIDDEN.into_response();
735        }
736        next.run(request).await
737    }
738
739    fn is_loopback(origin: &axum::http::HeaderValue) -> bool {
740        let Ok(origin) = origin.to_str() else {
741            return false;
742        };
743        let Ok(url) = Url::parse(origin) else {
744            return false;
745        };
746        match url.host() {
747            Some(Host::Domain("localhost")) => true,
748            Some(Host::Ipv4(address)) => address.is_loopback(),
749            Some(Host::Ipv6(address)) => address.is_loopback(),
750            _ => false,
751        }
752    }
753}
tools_kit/sources/mcp/server/mod.rs

tools_kit/sources/mcp/server/
mod.rs
