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

tools_kit/sources/mcp/server/
mod.rs
