fastmcp_server/
extensions.rs

1//! Typed server handlers for negotiated MCP extensions.
2//!
3//! This registry deliberately owns neither JSON-RPC decoding nor server/router
4//! dispatch. It binds a typed handler to an extension capability and request
5//! method, then delegates current-exchange admission to the frozen protocol
6//! descriptor registry before invoking that handler. Router integration owns
7//! the wire boundary separately.
8
9use std::collections::BTreeMap;
10use std::fmt;
11use std::marker::PhantomData;
12
13use fastmcp_core::{McpContext, McpError, McpResult};
14use fastmcp_protocol::extensions::{
15    ExtensionDispatchError, ExtensionRegistryError, MAX_EXTENSION_MEMBER_NAME_BYTES,
16    NegotiatedExtensionSet, ServerExtensionDiscovery,
17};
18#[cfg(feature = "apps")]
19use fastmcp_protocol::extensions::{
20    official_mcp_apps_empty_server_settings, register_official_mcp_apps_extension,
21    validate_official_mcp_apps_server_settings,
22};
23use fastmcp_protocol::protocol_policy::ProtocolEra;
24use fastmcp_protocol::{
25    ExtensionDescriptorRegistry, ExtensionDirection, ExtensionId, ExtensionRegistryReceipt,
26    ExtensionSettings, JsonRpcRequest,
27};
28use serde::Serialize;
29use serde::de::DeserializeOwned;
30use serde_json::Value;
31
32/// A typed implementation for one client-to-server extension request.
33///
34/// The supplied [`McpContext`] is the caller's request context. Implementations
35/// must preserve its cancellation and budget ownership rather than creating a
36/// detached runtime context.
37pub trait ExtensionHandler<Request, Response>: Send + Sync {
38    /// Handles one already-admitted extension request.
39    fn handle(&self, context: &McpContext, request: Request) -> McpResult<Response>;
40}
41
42#[cfg(all(test, feature = "apps"))]
43mod apps_only_tests {
44    use fastmcp_core::{McpContext, McpResult};
45    use fastmcp_protocol::extensions::official_mcp_apps_empty_server_settings;
46    use fastmcp_protocol::{
47        ExtensionDescriptorRegistry, ExtensionRegistryError, ExtensionSettings,
48    };
49    use serde_json::json;
50
51    use super::{ExtensionHandlerKey, ExtensionHandlerRegistrationError, ExtensionHandlerRegistry};
52
53    #[test]
54    fn official_apps_rejects_unowned_client_to_server_handler_without_mutation() {
55        let mut handlers = ExtensionHandlerRegistry::new(ExtensionDescriptorRegistry::new());
56        let apps_id = handlers
57            .install_official_mcp_apps()
58            .expect("the official Apps descriptor and marker install");
59        let key = ExtensionHandlerKey::new(
60            apps_id.clone(),
61            fastmcp_protocol::extensions::MCP_APPS_INITIALIZE_METHOD,
62        );
63
64        assert_eq!(
65            handlers.register(
66                apps_id.clone(),
67                fastmcp_protocol::extensions::MCP_APPS_INITIALIZE_METHOD,
68                |_context: &McpContext,
69                 _params: serde_json::Value|
70                 -> McpResult<serde_json::Value> { Ok(json!({})) },
71            ),
72            Err(ExtensionHandlerRegistrationError::MethodNotOwned(key)),
73            "MCP Apps owns no client-to-server extension method on this server registry"
74        );
75        assert_eq!(
76            handlers.len(),
77            0,
78            "the rejected Apps handler cannot create a dead dispatch entry"
79        );
80        assert_eq!(
81            handlers.server_metadata_len(),
82            1,
83            "the rejection preserves the installed Apps discovery marker"
84        );
85        assert_eq!(
86            handlers.descriptor_registry().descriptor(&apps_id),
87            Some(&fastmcp_protocol::official_mcp_apps_descriptor()),
88            "the rejected handler cannot alter the official Apps descriptor"
89        );
90    }
91
92    #[test]
93    fn official_apps_installation_duplicate_is_rejected_without_mutating_metadata() {
94        let mut handlers = ExtensionHandlerRegistry::new(ExtensionDescriptorRegistry::new());
95        let apps_id = handlers
96            .install_official_mcp_apps()
97            .expect("baseline official Apps installation succeeds");
98
99        assert_eq!(
100            handlers
101                .install_official_mcp_apps()
102                .expect_err("only the duplicate installation is rejected"),
103            ExtensionHandlerRegistrationError::OfficialMcpAppsAlreadyInstalled
104        );
105        assert_eq!(handlers.server_metadata_len(), 1);
106        assert_eq!(
107            handlers.descriptor_registry().descriptor(&apps_id),
108            Some(&fastmcp_protocol::official_mcp_apps_descriptor())
109        );
110    }
111
112    #[test]
113    fn manually_registered_official_apps_metadata_requires_the_empty_marker() {
114        let mut descriptors = ExtensionDescriptorRegistry::new();
115        let apps_id = fastmcp_protocol::register_official_mcp_apps_extension(&mut descriptors)
116            .expect("the official Apps descriptor registers");
117        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
118        let rejected = ExtensionSettings::new(json!({ "unexpected": true }))
119            .expect("the one-field alternate is generic extension metadata");
120
121        assert_eq!(
122            handlers.register_server_metadata(apps_id.clone(), rejected),
123            Err(ExtensionHandlerRegistrationError::Registry(
124                ExtensionRegistryError::OfficialMcpAppsServerSettingsNotEmpty
125            )),
126            "only the non-empty official Apps marker is rejected"
127        );
128        assert_eq!(
129            handlers.server_metadata_len(),
130            0,
131            "the rejected marker cannot be retained for later builder composition"
132        );
133
134        handlers
135            .register_server_metadata(apps_id, official_mcp_apps_empty_server_settings())
136            .expect("the exact empty official Apps marker remains accepted");
137        assert_eq!(handlers.server_metadata_len(), 1);
138    }
139}
140
141impl<Request, Response, Handler> ExtensionHandler<Request, Response> for Handler
142where
143    Handler: Fn(&McpContext, Request) -> McpResult<Response> + Send + Sync,
144{
145    fn handle(&self, context: &McpContext, request: Request) -> McpResult<Response> {
146        self(context, request)
147    }
148}
149
150/// Object-safe bridge that retains each registered handler's Rust types.
151trait ErasedExtensionHandler: Send + Sync {
152    fn invoke(&self, context: &McpContext, parameters: Value) -> McpResult<Value>;
153}
154
155/// Typed serde adapter retained behind the heterogeneous handler map.
156struct SerdeExtensionHandler<Request, Response, Handler> {
157    handler: Handler,
158    marker: PhantomData<fn(Request) -> Response>,
159}
160
161impl<Request, Response, Handler> ErasedExtensionHandler
162    for SerdeExtensionHandler<Request, Response, Handler>
163where
164    Request: DeserializeOwned,
165    Response: Serialize,
166    Handler: ExtensionHandler<Request, Response>,
167{
168    fn invoke(&self, context: &McpContext, parameters: Value) -> McpResult<Value> {
169        let request = serde_json::from_value(parameters)
170            .map_err(|error| McpError::invalid_params(error.to_string()))?;
171        let response = self.handler.handle(context, request)?;
172        serde_json::to_value(response).map_err(|_| {
173            McpError::internal_error("typed extension handler response serialization failed")
174        })
175    }
176}
177
178/// A stable server-side location for one extension request handler.
179#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
180pub struct ExtensionHandlerKey {
181    extension_id: ExtensionId,
182    method: String,
183}
184
185impl ExtensionHandlerKey {
186    /// Creates a handler location from its capability and exact request method.
187    #[must_use]
188    pub fn new(extension_id: ExtensionId, method: impl Into<String>) -> Self {
189        Self {
190            extension_id,
191            method: method.into(),
192        }
193    }
194
195    /// Returns the owning extension capability.
196    #[must_use]
197    pub const fn extension_id(&self) -> &ExtensionId {
198        &self.extension_id
199    }
200
201    /// Returns the exact client-to-server method spelling.
202    #[must_use]
203    pub fn method(&self) -> &str {
204        &self.method
205    }
206}
207
208/// Failure while registering a typed extension handler.
209#[derive(Clone, Debug, Eq, PartialEq)]
210pub enum ExtensionHandlerRegistrationError {
211    /// Handler mutation was attempted after the registry was frozen.
212    Frozen,
213    /// The handler's capability is absent from the protocol descriptor registry.
214    UnregisteredExtension(String),
215    /// The handler method is empty and cannot name an RPC request.
216    EmptyMethodName,
217    /// The handler method exceeds the protocol's bounded member-name limit.
218    MethodNameTooLong(String),
219    /// The extension descriptor does not own the requested method.
220    MethodNotOwned(ExtensionHandlerKey),
221    /// The extension descriptor owns the method only in the server-to-client direction.
222    MethodNotClientToServer(ExtensionHandlerKey),
223    /// A handler is already registered for this exact extension request location.
224    DuplicateHandler(ExtensionHandlerKey),
225    /// Server discovery metadata is already registered for this extension.
226    DuplicateServerMetadata(ExtensionId),
227    /// The official MCP Apps descriptor and server marker are already installed.
228    #[cfg(feature = "apps")]
229    OfficialMcpAppsAlreadyInstalled,
230    /// The protocol registry rejected official MCP Apps descriptor installation.
231    Registry(ExtensionRegistryError),
232}
233
234impl fmt::Display for ExtensionHandlerRegistrationError {
235    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
236        match self {
237            Self::Frozen => formatter.write_str("extension handler registry is frozen"),
238            Self::UnregisteredExtension(id) => {
239                write!(formatter, "extension handler has no descriptor: {id}")
240            }
241            Self::EmptyMethodName => formatter.write_str("extension handler method is empty"),
242            Self::MethodNameTooLong(method) => {
243                write!(
244                    formatter,
245                    "extension handler method exceeds its byte limit: {method}"
246                )
247            }
248            Self::MethodNotOwned(key) => write!(
249                formatter,
250                "extension handler method is not owned by its descriptor: {}/{}",
251                key.extension_id(),
252                key.method()
253            ),
254            Self::MethodNotClientToServer(key) => write!(
255                formatter,
256                "extension handler method is not client-to-server: {}/{}",
257                key.extension_id(),
258                key.method()
259            ),
260            Self::DuplicateHandler(key) => write!(
261                formatter,
262                "extension handler is already registered: {}/{}",
263                key.extension_id(),
264                key.method()
265            ),
266            Self::DuplicateServerMetadata(id) => {
267                write!(
268                    formatter,
269                    "extension server metadata is already registered: {id}"
270                )
271            }
272            #[cfg(feature = "apps")]
273            Self::OfficialMcpAppsAlreadyInstalled => {
274                formatter.write_str("official MCP Apps extension is already installed")
275            }
276            Self::Registry(error) => {
277                write!(formatter, "extension registry rejected install: {error}")
278            }
279        }
280    }
281}
282
283impl std::error::Error for ExtensionHandlerRegistrationError {
284    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
285        match self {
286            Self::Registry(error) => Some(error),
287            Self::Frozen
288            | Self::UnregisteredExtension(_)
289            | Self::EmptyMethodName
290            | Self::MethodNameTooLong(_)
291            | Self::MethodNotOwned(_)
292            | Self::MethodNotClientToServer(_)
293            | Self::DuplicateHandler(_)
294            | Self::DuplicateServerMetadata(_) => None,
295            #[cfg(feature = "apps")]
296            Self::OfficialMcpAppsAlreadyInstalled => None,
297        }
298    }
299}
300
301/// Failure while looking up a registered extension handler.
302#[derive(Clone, Debug, Eq, PartialEq)]
303pub enum ExtensionHandlerLookupError {
304    /// Lookup is unavailable until the capability and handler registries are frozen together.
305    RegistryNotFrozen,
306    /// No typed handler is registered for the admitted capability and request method.
307    HandlerNotFound(ExtensionHandlerKey),
308}
309
310impl fmt::Display for ExtensionHandlerLookupError {
311    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
312        match self {
313            Self::RegistryNotFrozen => {
314                formatter.write_str("extension handler registry must be frozen before lookup")
315            }
316            Self::HandlerNotFound(key) => write!(
317                formatter,
318                "no extension handler is registered: {}/{}",
319                key.extension_id(),
320                key.method()
321            ),
322        }
323    }
324}
325
326impl std::error::Error for ExtensionHandlerLookupError {}
327
328/// Failure while admitting or invoking a typed extension handler.
329#[derive(Debug)]
330pub enum ExtensionHandlerInvocationError {
331    /// Invocation is unavailable until the capability and handler registries are frozen together.
332    RegistryNotFrozen,
333    /// The current exchange did not admit this capability and client-to-server request method.
334    Protocol(ExtensionDispatchError),
335    /// No handler was registered for an otherwise admitted extension request method.
336    HandlerNotFound(ExtensionHandlerKey),
337    /// A descriptor-owned extension request was sent as an id-less notification.
338    RequestEnvelopeRequired(ExtensionHandlerKey),
339    /// The admitted typed handler returned an MCP error.
340    Handler(McpError),
341}
342
343impl fmt::Display for ExtensionHandlerInvocationError {
344    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
345        match self {
346            Self::RegistryNotFrozen => {
347                formatter.write_str("extension handler registry must be frozen before invocation")
348            }
349            Self::Protocol(error) => {
350                write!(formatter, "extension request was not admitted: {error}")
351            }
352            Self::HandlerNotFound(key) => write!(
353                formatter,
354                "no extension handler is registered: {}/{}",
355                key.extension_id(),
356                key.method()
357            ),
358            Self::RequestEnvelopeRequired(key) => write!(
359                formatter,
360                "extension request requires a JSON-RPC request id: {}/{}",
361                key.extension_id(),
362                key.method()
363            ),
364            Self::Handler(error) => write!(formatter, "extension handler failed: {error}"),
365        }
366    }
367}
368
369impl std::error::Error for ExtensionHandlerInvocationError {
370    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
371        match self {
372            Self::Protocol(error) => Some(error),
373            Self::Handler(error) => Some(error),
374            Self::RegistryNotFrozen
375            | Self::HandlerNotFound(_)
376            | Self::RequestEnvelopeRequired(_) => None,
377        }
378    }
379}
380
381/// Frozen, typed server handlers bound to one protocol descriptor registry.
382///
383/// The map is erased only internally. Each [`Self::register`] call fixes that
384/// method's own decoded request and encoded response types in a serde adapter,
385/// so one registry can hold heterogeneous extension methods without weakening
386/// a handler's typed Rust boundary. Method ownership, direction, negotiated
387/// capability activation, and descriptor receipt are still enforced by the
388/// protocol registry at every invocation.
389pub struct ExtensionHandlerRegistry {
390    descriptor_registry: ExtensionDescriptorRegistry,
391    handlers: BTreeMap<ExtensionHandlerKey, Box<dyn ErasedExtensionHandler>>,
392    server_metadata: BTreeMap<ExtensionId, ExtensionSettings>,
393    frozen: bool,
394}
395
396impl ExtensionHandlerRegistry {
397    /// Starts a handler registry around the supplied protocol descriptor registry.
398    #[must_use]
399    pub fn new(descriptor_registry: ExtensionDescriptorRegistry) -> Self {
400        Self {
401            descriptor_registry,
402            handlers: BTreeMap::new(),
403            server_metadata: BTreeMap::new(),
404            frozen: false,
405        }
406    }
407
408    /// Returns the protocol descriptor registry that governs handler admission.
409    #[must_use]
410    pub const fn descriptor_registry(&self) -> &ExtensionDescriptorRegistry {
411        &self.descriptor_registry
412    }
413
414    /// Returns the mutable descriptor registry before handler freeze.
415    ///
416    /// Builder-only composition must add every descriptor before typed handlers
417    /// are frozen under their shared receipt. Once frozen, neither descriptor
418    /// nor handler mutation is admitted.
419    pub(crate) fn descriptor_registry_mut(
420        &mut self,
421    ) -> Result<&mut ExtensionDescriptorRegistry, ExtensionHandlerRegistrationError> {
422        if self.frozen {
423            return Err(ExtensionHandlerRegistrationError::Frozen);
424        }
425        Ok(&mut self.descriptor_registry)
426    }
427
428    /// Returns whether handler registration has been frozen.
429    #[must_use]
430    pub const fn is_frozen(&self) -> bool {
431        self.frozen
432    }
433
434    /// Returns the number of registered typed handlers.
435    #[must_use]
436    pub fn len(&self) -> usize {
437        self.handlers.len()
438    }
439
440    /// Returns whether no typed handlers have been registered.
441    #[must_use]
442    pub fn is_empty(&self) -> bool {
443        self.handlers.is_empty()
444    }
445
446    /// Returns the number of descriptor-bound server metadata entries.
447    #[must_use]
448    pub fn server_metadata_len(&self) -> usize {
449        self.server_metadata.len()
450    }
451
452    /// Returns server metadata registered alongside descriptors before freeze.
453    ///
454    /// The live server runtime merges this snapshot with caller-supplied
455    /// discovery before it derives local extension enablement. The public
456    /// [`Self::server_discovery`] surface remains frozen-only.
457    pub(crate) fn configured_server_discovery(&self) -> ServerExtensionDiscovery {
458        ServerExtensionDiscovery {
459            extensions: self.server_metadata.clone(),
460        }
461    }
462
463    /// Registers the exact server discovery settings for one extension.
464    ///
465    /// This is intentionally separate from request-handler registration: an
466    /// extension such as MCP Apps can need a server discovery marker before it
467    /// owns a client-to-server extension method. The marker is accepted only
468    /// for a descriptor already linked into this registry, and is immutable
469    /// once the descriptor receipt is frozen.
470    pub fn register_server_metadata(
471        &mut self,
472        extension_id: ExtensionId,
473        settings: ExtensionSettings,
474    ) -> Result<(), ExtensionHandlerRegistrationError> {
475        if self.frozen {
476            return Err(ExtensionHandlerRegistrationError::Frozen);
477        }
478        if self.descriptor_registry.descriptor(&extension_id).is_none() {
479            return Err(ExtensionHandlerRegistrationError::UnregisteredExtension(
480                extension_id.to_string(),
481            ));
482        }
483        #[cfg(feature = "apps")]
484        if extension_id == fastmcp_protocol::official_mcp_apps_extension_id() {
485            validate_official_mcp_apps_server_settings(&settings)
486                .map_err(ExtensionHandlerRegistrationError::Registry)?;
487        }
488        if self.server_metadata.contains_key(&extension_id) {
489            return Err(ExtensionHandlerRegistrationError::DuplicateServerMetadata(
490                extension_id,
491            ));
492        }
493
494        self.server_metadata.insert(extension_id, settings);
495        Ok(())
496    }
497
498    /// Installs the official MCP Apps descriptor and its exact empty server marker.
499    ///
500    /// MCP Apps has no client-to-server extension method, so its server-side
501    /// installation consists solely of the protocol descriptor and discovery
502    /// metadata. The resolver remains caller-owned and must preserve its
503    /// `ExtensionSettingsCompatibilityResolver::resolve_with_disposition`
504    /// result when this registry is installed into a live server runtime.
505    #[cfg(feature = "apps")]
506    pub fn install_official_mcp_apps(
507        &mut self,
508    ) -> Result<ExtensionId, ExtensionHandlerRegistrationError> {
509        if self.frozen {
510            return Err(ExtensionHandlerRegistrationError::Frozen);
511        }
512        let id = fastmcp_protocol::official_mcp_apps_extension_id();
513        if self.descriptor_registry.descriptor(&id).is_some()
514            || self.server_metadata.contains_key(&id)
515        {
516            return Err(ExtensionHandlerRegistrationError::OfficialMcpAppsAlreadyInstalled);
517        }
518        register_official_mcp_apps_extension(&mut self.descriptor_registry)
519            .map_err(ExtensionHandlerRegistrationError::Registry)?;
520        self.register_server_metadata(id.clone(), official_mcp_apps_empty_server_settings())?;
521        Ok(id)
522    }
523
524    /// Registers one typed handler for an extension request method.
525    ///
526    /// Descriptor existence, exact method ownership, and client-to-server
527    /// direction are checked immediately. [`Self::invoke`] still checks
528    /// request-specific admission against the current negotiated extension set
529    /// before the handler can run.
530    pub fn register<Request, Response, Handler>(
531        &mut self,
532        extension_id: ExtensionId,
533        method: impl Into<String>,
534        handler: Handler,
535    ) -> Result<(), ExtensionHandlerRegistrationError>
536    where
537        Request: DeserializeOwned + 'static,
538        Response: Serialize + 'static,
539        Handler: ExtensionHandler<Request, Response> + 'static,
540    {
541        if self.frozen {
542            return Err(ExtensionHandlerRegistrationError::Frozen);
543        }
544        if self.descriptor_registry.descriptor(&extension_id).is_none() {
545            return Err(ExtensionHandlerRegistrationError::UnregisteredExtension(
546                extension_id.to_string(),
547            ));
548        }
549
550        let key = ExtensionHandlerKey::new(extension_id, method);
551        if key.method().is_empty() {
552            return Err(ExtensionHandlerRegistrationError::EmptyMethodName);
553        }
554        if key.method().len() > MAX_EXTENSION_MEMBER_NAME_BYTES {
555            return Err(ExtensionHandlerRegistrationError::MethodNameTooLong(
556                key.method().to_owned(),
557            ));
558        }
559        let Some(method_descriptor) = self
560            .descriptor_registry
561            .method_descriptor(key.extension_id(), key.method())
562        else {
563            return Err(ExtensionHandlerRegistrationError::MethodNotOwned(key));
564        };
565        if method_descriptor.direction != ExtensionDirection::ClientToServer {
566            return Err(ExtensionHandlerRegistrationError::MethodNotClientToServer(
567                key,
568            ));
569        }
570        if self.handlers.contains_key(&key) {
571            return Err(ExtensionHandlerRegistrationError::DuplicateHandler(key));
572        }
573
574        self.handlers.insert(
575            key,
576            Box::new(SerdeExtensionHandler::<Request, Response, Handler> {
577                handler,
578                marker: PhantomData,
579            }),
580        );
581        Ok(())
582    }
583
584    /// Freezes protocol descriptors and typed handlers under one immutable receipt.
585    pub fn freeze(&mut self) -> Result<ExtensionRegistryReceipt, ExtensionRegistryError> {
586        let receipt = self.descriptor_registry.freeze()?;
587        self.frozen = true;
588        Ok(receipt)
589    }
590
591    /// Freezes this registry and returns its matching discovery metadata.
592    ///
593    /// Use this after extension composition is complete to inspect one frozen
594    /// descriptor receipt with its matching metadata. During builder
595    /// composition, [`ServerBuilder::extension_registry`](crate::ServerBuilder::extension_registry)
596    /// derives registered metadata directly so later official Tasks installation
597    /// remains mutable.
598    pub fn freeze_with_server_discovery(
599        &mut self,
600    ) -> Result<(ExtensionRegistryReceipt, ServerExtensionDiscovery), ExtensionRegistryError> {
601        let receipt = self.freeze()?;
602        let discovery = self
603            .server_discovery()
604            .expect("registry is frozen immediately after a successful freeze");
605        Ok((receipt, discovery))
606    }
607
608    /// Returns the server discovery data bound to this frozen registry.
609    ///
610    /// This frozen diagnostic/export surface is separate from builder-time
611    /// composition, which snapshots registered metadata before freeze.
612    pub fn server_discovery(
613        &self,
614    ) -> Result<ServerExtensionDiscovery, ExtensionHandlerLookupError> {
615        if !self.frozen {
616            return Err(ExtensionHandlerLookupError::RegistryNotFrozen);
617        }
618        Ok(ServerExtensionDiscovery {
619            extensions: self.server_metadata.clone(),
620        })
621    }
622
623    /// Looks up one handler after freezing has made the registry immutable.
624    pub fn lookup(
625        &self,
626        extension_id: &ExtensionId,
627        method: &str,
628    ) -> Result<&ExtensionHandlerKey, ExtensionHandlerLookupError> {
629        if !self.frozen {
630            return Err(ExtensionHandlerLookupError::RegistryNotFrozen);
631        }
632        let key = ExtensionHandlerKey::new(extension_id.clone(), method);
633        self.handlers
634            .get_key_value(&key)
635            .map(|(registered_key, _)| registered_key)
636            .ok_or(ExtensionHandlerLookupError::HandlerNotFound(key))
637    }
638
639    /// Admits and invokes one typed client-to-server extension request.
640    ///
641    /// The protocol admission call binds this invocation to the exact frozen
642    /// descriptor receipt, current negotiated extension set, modern era, and
643    /// client-to-server ownership and request envelope before the handler is
644    /// looked up or executed. A handler can only be registered for a
645    /// descriptor's request method, so an id-less notification is rejected
646    /// generically before its parameters can reach stateful extension code.
647    pub fn invoke(
648        &self,
649        negotiated: &NegotiatedExtensionSet,
650        protocol_era: ProtocolEra,
651        extension_id: &ExtensionId,
652        request: &JsonRpcRequest,
653        context: &McpContext,
654    ) -> Result<Value, ExtensionHandlerInvocationError> {
655        if !self.frozen {
656            return Err(ExtensionHandlerInvocationError::RegistryNotFrozen);
657        }
658
659        negotiated
660            .admit_method(
661                &self.descriptor_registry,
662                protocol_era,
663                extension_id,
664                &request.method,
665                ExtensionDirection::ClientToServer,
666            )
667            .map_err(ExtensionHandlerInvocationError::Protocol)?;
668        let key = ExtensionHandlerKey::new(extension_id.clone(), &request.method);
669        if request.is_notification() {
670            return Err(ExtensionHandlerInvocationError::RequestEnvelopeRequired(
671                key,
672            ));
673        }
674        let key = self
675            .lookup(extension_id, &request.method)
676            .map_err(|error| match error {
677                ExtensionHandlerLookupError::RegistryNotFrozen => {
678                    ExtensionHandlerInvocationError::RegistryNotFrozen
679                }
680                ExtensionHandlerLookupError::HandlerNotFound(key) => {
681                    ExtensionHandlerInvocationError::HandlerNotFound(key)
682                }
683            })?;
684        let handler = self
685            .handlers
686            .get(key)
687            .ok_or_else(|| ExtensionHandlerInvocationError::HandlerNotFound(key.clone()))?;
688        handler
689            .invoke(
690                context,
691                request
692                    .params
693                    .clone()
694                    .unwrap_or_else(|| serde_json::json!({})),
695            )
696            .map_err(ExtensionHandlerInvocationError::Handler)
697    }
698}
699
700#[cfg(all(test, feature = "tasks"))]
701mod tests {
702    use std::collections::BTreeMap;
703    use std::sync::Arc;
704    use std::sync::atomic::{AtomicUsize, Ordering};
705
706    use asupersync::Cx;
707    use fastmcp_core::{McpContext, McpErrorCode, McpResult};
708    #[cfg(feature = "apps")]
709    use fastmcp_protocol::extensions::official_mcp_apps_empty_server_settings;
710    use fastmcp_protocol::extensions::{
711        ClientExtensionDiscovery, ExtensionLocalEnablement, ExtensionSettings,
712        ServerExtensionDiscovery, official_tasks_descriptor, official_tasks_empty_settings,
713        register_official_tasks_extension,
714    };
715    use fastmcp_protocol::protocol_policy::ProtocolEra;
716    use fastmcp_protocol::{
717        ExtensionDescriptorRegistry, ExtensionDirection, ExtensionId, JsonRpcRequest,
718    };
719    use serde::{Deserialize, Serialize};
720    use serde_json::json;
721
722    use super::{
723        ExtensionHandlerInvocationError, ExtensionHandlerKey, ExtensionHandlerRegistrationError,
724        ExtensionHandlerRegistry,
725    };
726
727    fn tasks_descriptors() -> (ExtensionDescriptorRegistry, ExtensionId) {
728        let mut descriptors = ExtensionDescriptorRegistry::new();
729        let id = register_official_tasks_extension(&mut descriptors)
730            .expect("official Tasks descriptor registers");
731        (descriptors, id)
732    }
733
734    fn primary_tasks_descriptor_with_direction(
735        direction: ExtensionDirection,
736    ) -> (ExtensionDescriptorRegistry, ExtensionId) {
737        let id = fastmcp_protocol::official_tasks_extension_id();
738        let mut descriptor = official_tasks_descriptor();
739        let method = descriptor
740            .method
741            .as_mut()
742            .expect("official Tasks descriptor owns a primary request method");
743        method.direction = direction;
744        if direction == ExtensionDirection::ServerToClient {
745            method.http_era_disposition = None;
746        }
747        let mut descriptors = ExtensionDescriptorRegistry::new();
748        descriptors
749            .register(descriptor)
750            .expect("direction-adjusted Tasks descriptor remains structurally valid");
751        (descriptors, id)
752    }
753
754    #[derive(Deserialize)]
755    struct GetTaskRequest {
756        value: u32,
757    }
758
759    #[derive(Debug, PartialEq, Serialize)]
760    struct GetTaskResponse {
761        next: u32,
762    }
763
764    #[derive(Deserialize)]
765    struct UpdateTaskRequest {
766        title: String,
767    }
768
769    #[derive(Debug, PartialEq, Serialize)]
770    struct UpdateTaskResponse {
771        updated_title: String,
772    }
773
774    #[allow(
775        clippy::unnecessary_wraps,
776        reason = "the ExtensionHandler callback contract requires McpResult"
777    )]
778    fn get_task(context: &McpContext, request: GetTaskRequest) -> McpResult<GetTaskResponse> {
779        assert_eq!(context.request_id(), 71);
780        Ok(GetTaskResponse {
781            next: request.value + 1,
782        })
783    }
784
785    #[allow(
786        clippy::unnecessary_wraps,
787        reason = "the ExtensionHandler callback contract requires McpResult"
788    )]
789    fn alternate_get_task(
790        context: &McpContext,
791        request: GetTaskRequest,
792    ) -> McpResult<GetTaskResponse> {
793        assert_eq!(context.request_id(), 71);
794        Ok(GetTaskResponse {
795            next: request.value + 2,
796        })
797    }
798
799    #[allow(
800        clippy::unnecessary_wraps,
801        reason = "the ExtensionHandler callback contract requires McpResult"
802    )]
803    fn update_task(
804        context: &McpContext,
805        request: UpdateTaskRequest,
806    ) -> McpResult<UpdateTaskResponse> {
807        assert_eq!(context.request_id(), 71);
808        Ok(UpdateTaskResponse {
809            updated_title: request.title.to_uppercase(),
810        })
811    }
812
813    fn negotiated_tasks(
814        descriptors: &ExtensionDescriptorRegistry,
815        id: &ExtensionId,
816    ) -> fastmcp_protocol::extensions::NegotiatedExtensionSet {
817        let client = ClientExtensionDiscovery {
818            extensions: BTreeMap::from([(id.clone(), official_tasks_empty_settings())]),
819        };
820        let server = ServerExtensionDiscovery {
821            extensions: BTreeMap::from([(id.clone(), official_tasks_empty_settings())]),
822        };
823        let mut local = ExtensionLocalEnablement::default();
824        local.enable(id.clone());
825        let mut resolver =
826            |_descriptor: &fastmcp_protocol::ExtensionDescriptor,
827             _client: &ExtensionSettings,
828             _server: &ExtensionSettings| { Ok(official_tasks_empty_settings()) };
829
830        descriptors
831            .negotiate(
832                ProtocolEra::Modern2026,
833                &local,
834                &client,
835                &server,
836                &mut resolver,
837            )
838            .expect("bilaterally advertised official Tasks negotiates")
839    }
840
841    #[test]
842    fn extension_handler_registry_freezes_and_looks_up_registered_handler() {
843        let (descriptors, id) = tasks_descriptors();
844        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
845        handlers
846            .register(id.clone(), "tasks/get", get_task)
847            .expect("typed Tasks handler registers");
848
849        let receipt = handlers.freeze().expect("handler registry freezes");
850
851        assert!(handlers.is_frozen());
852        assert_eq!(handlers.len(), 1);
853        assert_eq!(handlers.descriptor_registry().receipt(), Some(&receipt));
854        let registered = handlers
855            .lookup(&id, "tasks/get")
856            .expect("frozen registry finds its handler registration");
857        assert_eq!(registered.extension_id(), &id);
858        assert_eq!(registered.method(), "tasks/get");
859    }
860
861    #[test]
862    fn extension_handler_registry_requires_client_to_server_ownership_without_mutation() {
863        let (valid_descriptors, valid_id) =
864            primary_tasks_descriptor_with_direction(ExtensionDirection::ClientToServer);
865        let mut valid_handlers = ExtensionHandlerRegistry::new(valid_descriptors);
866        valid_handlers
867            .register(valid_id, "tasks/get", get_task)
868            .expect("a descriptor-owned client-to-server method registers");
869        assert_eq!(valid_handlers.len(), 1);
870
871        let (rejected_descriptors, rejected_id) =
872            primary_tasks_descriptor_with_direction(ExtensionDirection::ServerToClient);
873        let mut rejected_handlers = ExtensionHandlerRegistry::new(rejected_descriptors);
874        let key = ExtensionHandlerKey::new(rejected_id.clone(), "tasks/get");
875
876        assert_eq!(
877            rejected_handlers.register(rejected_id.clone(), "tasks/get", get_task),
878            Err(ExtensionHandlerRegistrationError::MethodNotClientToServer(
879                key
880            )),
881            "the structurally valid server-to-client direction must reject before insertion"
882        );
883        assert_eq!(
884            rejected_handlers.len(),
885            0,
886            "a server-to-client method cannot create a dead server handler entry"
887        );
888        assert_eq!(rejected_handlers.server_metadata_len(), 0);
889        assert_eq!(
890            rejected_handlers
891                .descriptor_registry()
892                .method_descriptor(&rejected_id, "tasks/get")
893                .map(|method| method.direction),
894            Some(ExtensionDirection::ServerToClient),
895            "rejection must not alter the owned descriptor"
896        );
897    }
898
899    #[test]
900    #[cfg(feature = "apps")]
901    fn official_apps_server_metadata_is_emitted_from_the_frozen_registry() {
902        let (descriptors, tasks_id) = tasks_descriptors();
903        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
904        let apps_id = handlers
905            .install_official_mcp_apps()
906            .expect("the official Apps descriptor and marker install alongside Tasks");
907        let (receipt, discovery) = handlers
908            .freeze_with_server_discovery()
909            .expect("Apps metadata registry freezes with matching discovery");
910
911        assert_eq!(handlers.descriptor_registry().receipt(), Some(&receipt));
912        assert_eq!(handlers.server_metadata_len(), 1);
913        assert!(
914            handlers
915                .descriptor_registry()
916                .descriptor(&tasks_id)
917                .is_some(),
918            "Apps installation preserves the already-registered Tasks descriptor"
919        );
920        assert_eq!(
921            handlers.descriptor_registry().descriptor(&apps_id),
922            Some(&fastmcp_protocol::official_mcp_apps_descriptor()),
923            "server discovery must use the protocol's official Apps descriptor"
924        );
925        assert_eq!(
926            discovery
927                .extensions
928                .get(&apps_id)
929                .map(|settings| serde_json::Value::Object(settings.as_object().clone())),
930            Some(json!({})),
931            "the official Apps server marker is emitted only through the frozen descriptor registry"
932        );
933    }
934
935    #[test]
936    fn unnegotiated_extension_metadata_is_rejected_without_mutating_the_registry() {
937        let (descriptors, _) = tasks_descriptors();
938        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
939        let unnegotiated = fastmcp_protocol::official_mcp_apps_extension_id();
940        let marker = ExtensionSettings::new(json!({}))
941            .expect("the unchanged official Apps server marker is an object");
942
943        assert_eq!(
944            handlers.register_server_metadata(unnegotiated.clone(), marker),
945            Err(ExtensionHandlerRegistrationError::UnregisteredExtension(
946                unnegotiated.to_string()
947            )),
948            "only the extension identifier differs from the registered Apps metadata path"
949        );
950        assert_eq!(handlers.server_metadata_len(), 0);
951        assert!(
952            handlers
953                .descriptor_registry()
954                .descriptor(&unnegotiated)
955                .is_none()
956        );
957
958        handlers.freeze().expect("unchanged registry still freezes");
959        assert!(
960            handlers
961                .server_discovery()
962                .expect("unchanged frozen registry still exports discovery")
963                .extensions
964                .is_empty(),
965            "the rejected unnegotiated extension cannot alter advertised metadata"
966        );
967    }
968
969    #[test]
970    fn extension_handler_registry_rejects_duplicate_key_one_variable_negative() {
971        let (descriptors, id) = tasks_descriptors();
972        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
973        handlers
974            .register(id.clone(), "tasks/get", get_task)
975            .expect("baseline handler registers");
976
977        assert_eq!(
978            handlers.register(id.clone(), "tasks/get", alternate_get_task),
979            Err(ExtensionHandlerRegistrationError::DuplicateHandler(
980                ExtensionHandlerKey::new(id, "tasks/get")
981            )),
982            "only the handler implementation changes from the registered request location"
983        );
984    }
985
986    #[test]
987    fn extension_handler_registry_erases_heterogeneous_types_and_rejects_malformed_input() {
988        let (descriptors, id) = tasks_descriptors();
989        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
990        let get_calls = Arc::new(AtomicUsize::new(0));
991        let counted_get_calls = Arc::clone(&get_calls);
992        handlers
993            .register(
994                id.clone(),
995                "tasks/get",
996                move |context: &McpContext, request: GetTaskRequest| {
997                    counted_get_calls.fetch_add(1, Ordering::Relaxed);
998                    get_task(context, request)
999                },
1000            )
1001            .expect("typed Tasks get handler registers");
1002        handlers
1003            .register(id.clone(), "tasks/update", update_task)
1004            .expect("differently typed Tasks update handler registers");
1005        handlers.freeze().expect("handler registry freezes");
1006        let negotiated = negotiated_tasks(handlers.descriptor_registry(), &id);
1007        let context = McpContext::new(Cx::for_testing(), 71);
1008
1009        assert_eq!(
1010            handlers
1011                .invoke(
1012                    &negotiated,
1013                    ProtocolEra::Modern2026,
1014                    &id,
1015                    &JsonRpcRequest::new("tasks/get", Some(json!({"value": 41})), 1_i64),
1016                    &context,
1017                )
1018                .expect("negotiated protocol admission invokes the typed get handler"),
1019            json!({"next": 42})
1020        );
1021        assert_eq!(
1022            handlers
1023                .invoke(
1024                    &negotiated,
1025                    ProtocolEra::Modern2026,
1026                    &id,
1027                    &JsonRpcRequest::new("tasks/update", Some(json!({"title": "review"})), 2_i64),
1028                    &context,
1029                )
1030                .expect("the same registry invokes the differently typed update handler"),
1031            json!({"updated_title": "REVIEW"})
1032        );
1033        assert_eq!(get_calls.load(Ordering::Relaxed), 1);
1034
1035        let error = handlers
1036            .invoke(
1037                &negotiated,
1038                ProtocolEra::Modern2026,
1039                &id,
1040                &JsonRpcRequest::new("tasks/get", Some(json!({"unexpected": true})), 3_i64),
1041                &context,
1042            )
1043            .expect_err("one malformed request field shape must reject before the handler runs");
1044        let ExtensionHandlerInvocationError::Handler(error) = error else {
1045            panic!("malformed parameters must reach the typed serde admission boundary");
1046        };
1047        assert_eq!(error.code, McpErrorCode::InvalidParams);
1048        assert_eq!(
1049            get_calls.load(Ordering::Relaxed),
1050            1,
1051            "failed typed decoding cannot invoke the registered handler"
1052        );
1053    }
1054
1055    #[test]
1056    fn extension_request_descriptor_rejects_an_idless_near_match_before_handler_invocation() {
1057        let (descriptors, id) = tasks_descriptors();
1058        let mut handlers = ExtensionHandlerRegistry::new(descriptors);
1059        let calls = Arc::new(AtomicUsize::new(0));
1060        let counted_calls = Arc::clone(&calls);
1061        handlers
1062            .register(
1063                id.clone(),
1064                "tasks/update",
1065                move |_context: &McpContext, _request: UpdateTaskRequest| {
1066                    counted_calls.fetch_add(1, Ordering::Relaxed);
1067                    Ok(serde_json::json!({"resultType": "complete"}))
1068                },
1069            )
1070            .expect("typed Tasks update handler registers");
1071        handlers.freeze().expect("handler registry freezes");
1072        let negotiated = negotiated_tasks(handlers.descriptor_registry(), &id);
1073        let context = McpContext::new(Cx::for_testing(), 72);
1074        let parameters = json!({"title": "review"});
1075
1076        assert_eq!(
1077            handlers
1078                .invoke(
1079                    &negotiated,
1080                    ProtocolEra::Modern2026,
1081                    &id,
1082                    &JsonRpcRequest::new("tasks/update", Some(parameters.clone()), 4_i64),
1083                    &context,
1084                )
1085                .expect("the request-shaped extension frame is admitted"),
1086            json!({"resultType": "complete"})
1087        );
1088        assert_eq!(calls.load(Ordering::Relaxed), 1);
1089
1090        let error = handlers
1091            .invoke(
1092                &negotiated,
1093                ProtocolEra::Modern2026,
1094                &id,
1095                &JsonRpcRequest::notification("tasks/update", Some(parameters)),
1096                &context,
1097            )
1098            .expect_err("removing only the request id must reject before handler invocation");
1099        assert!(matches!(
1100            error,
1101            ExtensionHandlerInvocationError::RequestEnvelopeRequired(_)
1102        ));
1103        assert_eq!(
1104            calls.load(Ordering::Relaxed),
1105            1,
1106            "the rejected id-less near-match must leave handler-owned state unchanged"
1107        );
1108    }
1109}
fastmcp_server/extensions.rs

fastmcp_server/
extensions.rs
