Skip to main content

tower_mcp/
router.rs

1//! MCP Router - routes requests to tools, resources, and prompts
2//!
3//! The router implements Tower's `Service` trait, making it composable with
4//! standard tower middleware.
5
6use std::collections::{HashMap, HashSet};
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::{Arc, RwLock};
10use std::task::{Context, Poll};
11
12use tower_service::Service;
13
14use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
15
16use crate::async_task::TaskStore;
17use crate::context::{
18    CancellationToken, ClientRequesterHandle, NotificationSender, RequestContext,
19    ServerNotification,
20};
21use crate::error::{Error, JsonRpcError, Result};
22use crate::filter::{PromptFilter, ResourceFilter, ToolFilter};
23use crate::prompt::Prompt;
24use crate::protocol::*;
25#[cfg(feature = "dynamic-tools")]
26use crate::registry::{
27    DynamicPromptRegistry, DynamicPromptsInner, DynamicResourceRegistry,
28    DynamicResourceTemplateRegistry, DynamicResourceTemplatesInner, DynamicResourcesInner,
29    DynamicToolRegistry, DynamicToolsInner,
30};
31use crate::resource::{Resource, ResourceTemplate};
32use crate::session::SessionState;
33use crate::tool::Tool;
34
35/// Type alias for completion handler function
36pub(crate) type CompletionHandler = Arc<
37    dyn Fn(CompleteParams) -> Pin<Box<dyn Future<Output = Result<CompleteResult>> + Send>>
38        + Send
39        + Sync,
40>;
41
42/// Decode a pagination cursor into an offset.
43///
44/// Returns `Err` if the cursor is malformed.
45fn decode_cursor(cursor: &str) -> Result<usize> {
46    let bytes = BASE64
47        .decode(cursor)
48        .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
49    let s = String::from_utf8(bytes)
50        .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))?;
51    s.parse::<usize>()
52        .map_err(|_| Error::JsonRpc(JsonRpcError::invalid_params("Invalid pagination cursor")))
53}
54
55/// Encode an offset into an opaque pagination cursor.
56fn encode_cursor(offset: usize) -> String {
57    BASE64.encode(offset.to_string())
58}
59
60/// Apply pagination to a collected list of items.
61///
62/// Returns the page of items and an optional `next_cursor`.
63fn paginate<T>(
64    items: Vec<T>,
65    cursor: Option<&str>,
66    page_size: Option<usize>,
67) -> Result<(Vec<T>, Option<String>)> {
68    let Some(page_size) = page_size else {
69        return Ok((items, None));
70    };
71
72    let offset = match cursor {
73        Some(c) => decode_cursor(c)?,
74        None => 0,
75    };
76
77    if offset >= items.len() {
78        return Ok((Vec::new(), None));
79    }
80
81    let end = (offset + page_size).min(items.len());
82    let next_cursor = if end < items.len() {
83        Some(encode_cursor(end))
84    } else {
85        None
86    };
87
88    let mut items = items;
89    let page = items.drain(offset..end).collect();
90    Ok((page, next_cursor))
91}
92
93/// MCP Router that dispatches requests to registered handlers
94///
95/// Implements `tower::Service<McpRequest>` for middleware composition.
96///
97/// # Example
98///
99/// ```rust
100/// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
101/// use schemars::JsonSchema;
102/// use serde::Deserialize;
103///
104/// #[derive(Debug, Deserialize, JsonSchema)]
105/// struct Input { value: String }
106///
107/// let tool = ToolBuilder::new("echo")
108///     .description("Echo input")
109///     .handler(|i: Input| async move { Ok(CallToolResult::text(i.value)) })
110///     .build();
111///
112/// let router = McpRouter::new()
113///     .server_info("my-server", "1.0.0")
114///     .tool(tool);
115/// ```
116#[derive(Clone)]
117pub struct McpRouter {
118    inner: Arc<McpRouterInner>,
119    session: SessionState,
120}
121
122impl std::fmt::Debug for McpRouter {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("McpRouter")
125            .field("server_name", &self.inner.server_name)
126            .field("server_version", &self.inner.server_version)
127            .field("tools_count", &self.inner.tools.len())
128            .field("resources_count", &self.inner.resources.len())
129            .field("prompts_count", &self.inner.prompts.len())
130            .field("session_phase", &self.session.phase())
131            .finish()
132    }
133}
134
135/// Configuration for auto-generated instructions
136#[derive(Clone, Debug)]
137struct AutoInstructionsConfig {
138    prefix: Option<String>,
139    suffix: Option<String>,
140}
141
142/// Inner configuration that is shared across clones
143#[derive(Clone)]
144struct McpRouterInner {
145    server_name: String,
146    server_version: String,
147    /// Human-readable title for the server
148    server_title: Option<String>,
149    /// Description of the server
150    server_description: Option<String>,
151    /// Icons for the server
152    server_icons: Option<Vec<ToolIcon>>,
153    /// URL of the server's website
154    server_website_url: Option<String>,
155    instructions: Option<String>,
156    auto_instructions: Option<AutoInstructionsConfig>,
157    tools: HashMap<String, Arc<Tool>>,
158    resources: HashMap<String, Arc<Resource>>,
159    /// Resource templates for dynamic resource matching (keyed by uri_template)
160    resource_templates: Vec<Arc<ResourceTemplate>>,
161    prompts: HashMap<String, Arc<Prompt>>,
162    /// In-flight requests for cancellation tracking (shared across clones)
163    in_flight: Arc<RwLock<HashMap<RequestId, CancellationToken>>>,
164    /// Channel for sending notifications to connected clients
165    notification_tx: Option<NotificationSender>,
166    /// Handle for sending requests to the client (for sampling, etc.)
167    client_requester: Option<ClientRequesterHandle>,
168    /// Task store for async operations
169    task_store: TaskStore,
170    /// Subscribed resource URIs
171    subscriptions: Arc<RwLock<HashSet<String>>>,
172    /// Handler for completion requests
173    completion_handler: Option<CompletionHandler>,
174    /// Filter for tools based on session state
175    tool_filter: Option<ToolFilter>,
176    /// Filter for resources based on session state
177    resource_filter: Option<ResourceFilter>,
178    /// Filter for prompts based on session state
179    prompt_filter: Option<PromptFilter>,
180    /// Router-level extensions (for state and middleware data)
181    extensions: Arc<crate::context::Extensions>,
182    /// Minimum log level for filtering outgoing log notifications (set by client via logging/setLevel)
183    min_log_level: Arc<RwLock<LogLevel>>,
184    /// Page size for list method pagination (None = return all results)
185    page_size: Option<usize>,
186    /// TTL hint for list responses in milliseconds (SEP-2549).
187    /// When set, the value is returned as `ttlMs` in tools/list, resources/list,
188    /// and prompts/list responses so clients can cache the list.
189    list_ttl_ms: Option<u64>,
190    /// Deprecation info for the logging capability (SEP-2577).
191    /// When set, included in the `logging` capability in the initialize result.
192    logging_deprecated: Option<tower_mcp_types::protocol::DeprecationInfo>,
193    /// Names of tools that are currently disabled (hidden from list/call).
194    disabled_tools: Arc<RwLock<HashSet<String>>>,
195    /// URIs of resources that are currently disabled (hidden from list/read).
196    disabled_resources: Arc<RwLock<HashSet<String>>>,
197    /// Names of prompts that are currently disabled (hidden from list/get).
198    disabled_prompts: Arc<RwLock<HashSet<String>>>,
199    /// Dynamic tools registry for runtime tool (de)registration
200    #[cfg(feature = "dynamic-tools")]
201    dynamic_tools: Option<Arc<DynamicToolsInner>>,
202    /// Dynamic prompts registry for runtime prompt (de)registration
203    #[cfg(feature = "dynamic-tools")]
204    dynamic_prompts: Option<Arc<DynamicPromptsInner>>,
205    /// Dynamic resources registry for runtime resource (de)registration
206    #[cfg(feature = "dynamic-tools")]
207    dynamic_resources: Option<Arc<DynamicResourcesInner>>,
208    /// Dynamic resource templates registry for runtime template (de)registration
209    #[cfg(feature = "dynamic-tools")]
210    dynamic_resource_templates: Option<Arc<DynamicResourceTemplatesInner>>,
211}
212
213impl McpRouterInner {
214    /// Generate instructions text from registered tools, resources, and prompts.
215    fn generate_instructions(&self, config: &AutoInstructionsConfig) -> String {
216        let mut parts = Vec::new();
217
218        if let Some(prefix) = &config.prefix {
219            parts.push(prefix.clone());
220        }
221
222        // Tools section
223        if !self.tools.is_empty() {
224            let mut lines = vec!["## Tools".to_string(), String::new()];
225            let mut tools: Vec<_> = self.tools.values().collect();
226            tools.sort_by(|a, b| a.name.cmp(&b.name));
227            for tool in tools {
228                let desc = tool.description.as_deref().unwrap_or("No description");
229                let tags = annotation_tags(tool.annotations.as_ref());
230                if tags.is_empty() {
231                    lines.push(format!("- **{}**: {}", tool.name, desc));
232                } else {
233                    lines.push(format!("- **{}**: {} [{}]", tool.name, desc, tags));
234                }
235            }
236            parts.push(lines.join("\n"));
237        }
238
239        // Resources section
240        if !self.resources.is_empty() || !self.resource_templates.is_empty() {
241            let mut lines = vec!["## Resources".to_string(), String::new()];
242            let mut resources: Vec<_> = self.resources.values().collect();
243            resources.sort_by(|a, b| a.uri.cmp(&b.uri));
244            for resource in resources {
245                let desc = resource.description.as_deref().unwrap_or("No description");
246                lines.push(format!("- **{}**: {}", resource.uri, desc));
247            }
248            let mut templates: Vec<_> = self.resource_templates.iter().collect();
249            templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
250            for template in templates {
251                let desc = template.description.as_deref().unwrap_or("No description");
252                lines.push(format!("- **{}**: {}", template.uri_template, desc));
253            }
254            parts.push(lines.join("\n"));
255        }
256
257        // Prompts section
258        if !self.prompts.is_empty() {
259            let mut lines = vec!["## Prompts".to_string(), String::new()];
260            let mut prompts: Vec<_> = self.prompts.values().collect();
261            prompts.sort_by(|a, b| a.name.cmp(&b.name));
262            for prompt in prompts {
263                let desc = prompt.description.as_deref().unwrap_or("No description");
264                lines.push(format!("- **{}**: {}", prompt.name, desc));
265            }
266            parts.push(lines.join("\n"));
267        }
268
269        if let Some(suffix) = &config.suffix {
270            parts.push(suffix.clone());
271        }
272
273        parts.join("\n\n")
274    }
275}
276
277/// Build annotation tags like "read-only, idempotent" from tool annotations.
278///
279/// Only includes tags that differ from the MCP spec defaults
280/// (read-only=false, idempotent=false). The destructive and open-world
281/// hints are omitted because they match the default assumptions.
282fn annotation_tags(annotations: Option<&crate::protocol::ToolAnnotations>) -> String {
283    let Some(ann) = annotations else {
284        return String::new();
285    };
286    let mut tags = Vec::new();
287    if ann.is_read_only() {
288        tags.push("read-only");
289    }
290    if ann.is_idempotent() {
291        tags.push("idempotent");
292    }
293    tags.join(", ")
294}
295
296impl McpRouter {
297    /// Create a new MCP router
298    pub fn new() -> Self {
299        Self {
300            inner: Arc::new(McpRouterInner {
301                server_name: "tower-mcp".to_string(),
302                server_version: env!("CARGO_PKG_VERSION").to_string(),
303                server_title: None,
304                server_description: None,
305                server_icons: None,
306                server_website_url: None,
307                instructions: None,
308                auto_instructions: None,
309                tools: HashMap::new(),
310                resources: HashMap::new(),
311                resource_templates: Vec::new(),
312                prompts: HashMap::new(),
313                in_flight: Arc::new(RwLock::new(HashMap::new())),
314                notification_tx: None,
315                client_requester: None,
316                task_store: TaskStore::new(),
317                subscriptions: Arc::new(RwLock::new(HashSet::new())),
318                extensions: Arc::new(crate::context::Extensions::new()),
319                completion_handler: None,
320                tool_filter: None,
321                resource_filter: None,
322                prompt_filter: None,
323                min_log_level: Arc::new(RwLock::new(LogLevel::Debug)),
324                page_size: None,
325                list_ttl_ms: None,
326                logging_deprecated: None,
327                disabled_tools: Arc::new(RwLock::new(HashSet::new())),
328                disabled_resources: Arc::new(RwLock::new(HashSet::new())),
329                disabled_prompts: Arc::new(RwLock::new(HashSet::new())),
330                #[cfg(feature = "dynamic-tools")]
331                dynamic_tools: None,
332                #[cfg(feature = "dynamic-tools")]
333                dynamic_prompts: None,
334                #[cfg(feature = "dynamic-tools")]
335                dynamic_resources: None,
336                #[cfg(feature = "dynamic-tools")]
337                dynamic_resource_templates: None,
338            }),
339            session: SessionState::new(),
340        }
341    }
342
343    /// Create a clone with fresh session state.
344    ///
345    /// Use this when creating a new logical session (e.g., per HTTP connection).
346    /// The router configuration (tools, resources, prompts) is shared, but the
347    /// session state (phase, extensions) is independent.
348    ///
349    /// This is typically called by transports when establishing a new client session.
350    pub fn with_fresh_session(&self) -> Self {
351        Self {
352            inner: self.inner.clone(),
353            session: SessionState::new(),
354        }
355    }
356
357    /// Build a map of tool names to their annotations.
358    ///
359    /// The returned [`ToolAnnotationsMap`] includes annotations from all
360    /// currently registered tools (both static and dynamic). Tools without
361    /// annotations are omitted from the map.
362    ///
363    /// This is used internally by transports to inject annotations into
364    /// request extensions, but can also be called directly for custom
365    /// middleware setups.
366    pub fn tool_annotations_map(&self) -> ToolAnnotationsMap {
367        let disabled = self.inner.disabled_tools.read().unwrap();
368        let mut map = HashMap::new();
369        for (name, tool) in &self.inner.tools {
370            if disabled.contains(name) {
371                continue;
372            }
373            if let Some(annotations) = &tool.annotations {
374                map.insert(name.clone(), annotations.clone());
375            }
376        }
377        #[cfg(feature = "dynamic-tools")]
378        if let Some(dynamic) = &self.inner.dynamic_tools {
379            for tool in dynamic.list() {
380                if disabled.contains(&tool.name) {
381                    continue;
382                }
383                // Static tools take precedence
384                if !map.contains_key(&tool.name)
385                    && let Some(ref annotations) = tool.annotations
386                {
387                    map.insert(tool.name.clone(), annotations.clone());
388                }
389            }
390        }
391        ToolAnnotationsMap { map: Arc::new(map) }
392    }
393
394    /// Get access to the task store for async operations
395    pub fn task_store(&self) -> &TaskStore {
396        &self.inner.task_store
397    }
398
399    /// Enable dynamic tool registration and return a registry handle.
400    ///
401    /// The returned [`DynamicToolRegistry`] can be used to add and remove tools
402    /// at runtime. Dynamic tools are merged with static tools when handling
403    /// `tools/list` and `tools/call` requests. Static tools take precedence
404    /// over dynamic tools when names collide.
405    ///
406    /// # Example
407    ///
408    /// ```rust
409    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
410    /// use schemars::JsonSchema;
411    /// use serde::Deserialize;
412    ///
413    /// #[derive(Debug, Deserialize, JsonSchema)]
414    /// struct Input { value: String }
415    ///
416    /// let (router, registry) = McpRouter::new()
417    ///     .server_info("my-server", "1.0.0")
418    ///     .with_dynamic_tools();
419    ///
420    /// // Register a tool at runtime
421    /// let tool = ToolBuilder::new("echo")
422    ///     .description("Echo input")
423    ///     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
424    ///     .build();
425    ///
426    /// registry.register(tool);
427    /// ```
428    #[cfg(feature = "dynamic-tools")]
429    pub fn with_dynamic_tools(mut self) -> (Self, DynamicToolRegistry) {
430        let inner_dyn = Arc::new(DynamicToolsInner::new());
431        Arc::make_mut(&mut self.inner).dynamic_tools = Some(inner_dyn.clone());
432        (self, DynamicToolRegistry::new(inner_dyn))
433    }
434
435    /// Enable dynamic prompt registration and return a registry handle.
436    ///
437    /// The returned [`DynamicPromptRegistry`] can be used to add and remove
438    /// prompts at runtime. Dynamic prompts are merged with static prompts
439    /// when handling `prompts/list` and `prompts/get` requests. Static
440    /// prompts take precedence over dynamic prompts when names collide.
441    ///
442    /// # Example
443    ///
444    /// ```rust
445    /// use tower_mcp::{McpRouter, PromptBuilder};
446    ///
447    /// let (router, registry) = McpRouter::new()
448    ///     .server_info("my-server", "1.0.0")
449    ///     .with_dynamic_prompts();
450    ///
451    /// let prompt = PromptBuilder::new("greet")
452    ///     .description("Greet someone")
453    ///     .user_message("Hello!");
454    ///
455    /// registry.register(prompt);
456    /// ```
457    #[cfg(feature = "dynamic-tools")]
458    pub fn with_dynamic_prompts(mut self) -> (Self, DynamicPromptRegistry) {
459        let inner_dyn = Arc::new(DynamicPromptsInner::new());
460        Arc::make_mut(&mut self.inner).dynamic_prompts = Some(inner_dyn.clone());
461        (self, DynamicPromptRegistry::new(inner_dyn))
462    }
463
464    /// Enable dynamic resource registration and return a registry handle.
465    ///
466    /// The returned [`DynamicResourceRegistry`] can be used to add and remove
467    /// resources at runtime. Dynamic resources are merged with static resources
468    /// when handling `resources/list` and `resources/read` requests. Static
469    /// resources take precedence over dynamic resources when URIs collide.
470    ///
471    /// # Example
472    ///
473    /// ```rust
474    /// use tower_mcp::{McpRouter, ResourceBuilder};
475    ///
476    /// let (router, registry) = McpRouter::new()
477    ///     .server_info("my-server", "1.0.0")
478    ///     .with_dynamic_resources();
479    ///
480    /// let resource = ResourceBuilder::new("file:///data.json")
481    ///     .name("Data")
482    ///     .text(r#"{"key": "value"}"#);
483    ///
484    /// registry.register(resource);
485    /// ```
486    #[cfg(feature = "dynamic-tools")]
487    pub fn with_dynamic_resources(mut self) -> (Self, DynamicResourceRegistry) {
488        let inner_dyn = Arc::new(DynamicResourcesInner::new());
489        Arc::make_mut(&mut self.inner).dynamic_resources = Some(inner_dyn.clone());
490        (self, DynamicResourceRegistry::new(inner_dyn))
491    }
492
493    /// Enable dynamic resource template registration and return a registry handle.
494    ///
495    /// The returned [`DynamicResourceTemplateRegistry`] can be used to add and
496    /// remove resource templates at runtime. Dynamic templates are checked
497    /// after static templates when handling `resources/read` requests.
498    ///
499    /// # Example
500    ///
501    /// ```rust,ignore
502    /// use tower_mcp::{McpRouter, ResourceTemplateBuilder};
503    ///
504    /// let (router, registry) = McpRouter::new()
505    ///     .server_info("my-server", "1.0.0")
506    ///     .with_dynamic_resource_templates();
507    ///
508    /// let template = ResourceTemplateBuilder::new("db://tables/{table}")
509    ///     .name("Database Table")
510    ///     .handler(|uri, vars| async move { /* ... */ });
511    ///
512    /// registry.register(template);
513    /// ```
514    #[cfg(feature = "dynamic-tools")]
515    pub fn with_dynamic_resource_templates(mut self) -> (Self, DynamicResourceTemplateRegistry) {
516        let inner_dyn = Arc::new(DynamicResourceTemplatesInner::new());
517        Arc::make_mut(&mut self.inner).dynamic_resource_templates = Some(inner_dyn.clone());
518        (self, DynamicResourceTemplateRegistry::new(inner_dyn))
519    }
520
521    /// Set the notification sender for progress reporting
522    ///
523    /// This is typically called by the transport layer to receive notifications.
524    pub fn with_notification_sender(mut self, tx: NotificationSender) -> Self {
525        let inner = Arc::make_mut(&mut self.inner);
526        // Also register the sender with dynamic registries so they can
527        // broadcast list-changed notifications to this session.
528        #[cfg(feature = "dynamic-tools")]
529        if let Some(ref dynamic_tools) = inner.dynamic_tools {
530            dynamic_tools.add_notification_sender(tx.clone());
531        }
532        #[cfg(feature = "dynamic-tools")]
533        if let Some(ref dynamic_prompts) = inner.dynamic_prompts {
534            dynamic_prompts.add_notification_sender(tx.clone());
535        }
536        #[cfg(feature = "dynamic-tools")]
537        if let Some(ref dynamic_resources) = inner.dynamic_resources {
538            dynamic_resources.add_notification_sender(tx.clone());
539        }
540        #[cfg(feature = "dynamic-tools")]
541        if let Some(ref dynamic_resource_templates) = inner.dynamic_resource_templates {
542            dynamic_resource_templates.add_notification_sender(tx.clone());
543        }
544        inner.notification_tx = Some(tx);
545        self
546    }
547
548    /// Get the notification sender (if configured)
549    pub fn notification_sender(&self) -> Option<&NotificationSender> {
550        self.inner.notification_tx.as_ref()
551    }
552
553    /// Set the client requester for server-to-client requests (sampling, etc.)
554    ///
555    /// This is typically called by bidirectional transports (WebSocket, stdio)
556    /// to enable tool handlers to send requests to the client.
557    pub fn with_client_requester(mut self, requester: ClientRequesterHandle) -> Self {
558        Arc::make_mut(&mut self.inner).client_requester = Some(requester);
559        self
560    }
561
562    /// Get the client requester (if configured)
563    pub fn client_requester(&self) -> Option<&ClientRequesterHandle> {
564        self.inner.client_requester.as_ref()
565    }
566
567    /// Add router-level state that handlers can access via the `Extension<T>` extractor.
568    ///
569    /// This is the recommended way to share state across all tools, resources, and prompts
570    /// in a router. The state is available to handlers via the [`crate::extract::Extension`]
571    /// extractor.
572    ///
573    /// # Example
574    ///
575    /// ```rust
576    /// use std::sync::Arc;
577    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
578    /// use tower_mcp::extract::{Extension, Json};
579    /// use schemars::JsonSchema;
580    /// use serde::Deserialize;
581    ///
582    /// #[derive(Clone)]
583    /// struct AppState {
584    ///     db_url: String,
585    /// }
586    ///
587    /// #[derive(Deserialize, JsonSchema)]
588    /// struct QueryInput {
589    ///     sql: String,
590    /// }
591    ///
592    /// let state = Arc::new(AppState { db_url: "postgres://...".into() });
593    ///
594    /// // Tool extracts state via Extension<T>
595    /// let query_tool = ToolBuilder::new("query")
596    ///     .description("Run a database query")
597    ///     .extractor_handler(
598    ///         (),
599    ///         |Extension(state): Extension<Arc<AppState>>, Json(input): Json<QueryInput>| async move {
600    ///             Ok(CallToolResult::text(format!("Query on {}: {}", state.db_url, input.sql)))
601    ///         },
602    ///     )
603    ///     .build();
604    ///
605    /// let router = McpRouter::new()
606    ///     .with_state(state)  // State is now available to all handlers
607    ///     .tool(query_tool);
608    /// ```
609    pub fn with_state<T: Clone + Send + Sync + 'static>(mut self, state: T) -> Self {
610        let inner = Arc::make_mut(&mut self.inner);
611        Arc::make_mut(&mut inner.extensions).insert(state);
612        self
613    }
614
615    /// Add an extension value that handlers can access via the `Extension<T>` extractor.
616    ///
617    /// This is a more general form of `with_state()` for when you need multiple
618    /// typed values available to handlers.
619    pub fn with_extension<T: Clone + Send + Sync + 'static>(self, value: T) -> Self {
620        self.with_state(value)
621    }
622
623    /// Get the router's extensions.
624    pub fn extensions(&self) -> &crate::context::Extensions {
625        &self.inner.extensions
626    }
627
628    /// Create a request context for tracking a request
629    ///
630    /// This registers the request for cancellation tracking and sets up
631    /// progress reporting, client requests, and router extensions if configured.
632    pub fn create_context(
633        &self,
634        request_id: RequestId,
635        progress_token: Option<ProgressToken>,
636    ) -> RequestContext {
637        self.create_context_with_extensions(request_id, progress_token, &Extensions::new())
638    }
639
640    /// Internal: build a `RequestContext` and additionally merge per-request
641    /// extensions on top of the router's extensions. Used by [`Service::call`]
642    /// to thread `RouterRequest.extensions` (e.g. SEP-2575 per-request
643    /// `_meta`) through to handlers.
644    pub(crate) fn create_context_with_extensions(
645        &self,
646        request_id: RequestId,
647        progress_token: Option<ProgressToken>,
648        per_request: &Extensions,
649    ) -> RequestContext {
650        let ctx = RequestContext::new(request_id.clone());
651
652        // Set up progress token if provided
653        let ctx = if let Some(token) = progress_token {
654            ctx.with_progress_token(token)
655        } else {
656            ctx
657        };
658
659        // Set up notification sender if configured
660        let ctx = if let Some(tx) = &self.inner.notification_tx {
661            ctx.with_notification_sender(tx.clone())
662        } else {
663            ctx
664        };
665
666        // Set up client requester if configured (for sampling support)
667        let ctx = if let Some(requester) = &self.inner.client_requester {
668            ctx.with_client_requester(requester.clone())
669        } else {
670            ctx
671        };
672
673        // Start with router-level extensions, then layer per-request extensions
674        // on top so they win on type collision. with_state() data stays
675        // visible; per-request meta (SEP-2575) is now reachable too.
676        let mut merged = (*self.inner.extensions).clone();
677        merged.merge(per_request);
678        let ctx = ctx.with_extensions(Arc::new(merged));
679
680        // Set up log level filtering
681        let ctx = ctx.with_min_log_level(self.inner.min_log_level.clone());
682
683        // Register for cancellation tracking
684        let token = ctx.cancellation_token();
685        if let Ok(mut in_flight) = self.inner.in_flight.write() {
686            in_flight.insert(request_id, token);
687        }
688
689        ctx
690    }
691
692    /// Remove a request from tracking (called when request completes)
693    pub fn complete_request(&self, request_id: &RequestId) {
694        if let Ok(mut in_flight) = self.inner.in_flight.write() {
695            in_flight.remove(request_id);
696        }
697    }
698
699    /// Cancel a tracked request
700    fn cancel_request(&self, request_id: &RequestId) -> bool {
701        let Ok(in_flight) = self.inner.in_flight.read() else {
702            return false;
703        };
704        let Some(token) = in_flight.get(request_id) else {
705            return false;
706        };
707        token.cancel();
708        true
709    }
710
711    /// Set server info
712    pub fn server_info(mut self, name: impl Into<String>, version: impl Into<String>) -> Self {
713        let inner = Arc::make_mut(&mut self.inner);
714        inner.server_name = name.into();
715        inner.server_version = version.into();
716        self
717    }
718
719    /// Set the page size for list method pagination.
720    ///
721    /// When set, list methods (`tools/list`, `resources/list`, etc.) will return
722    /// at most `page_size` items per response, with a `next_cursor` for fetching
723    /// subsequent pages. When `None` (the default), all items are returned in a
724    /// single response.
725    pub fn page_size(mut self, size: usize) -> Self {
726        Arc::make_mut(&mut self.inner).page_size = Some(size);
727        self
728    }
729
730    /// Set a TTL hint on list responses (tools/list, resources/list, prompts/list).
731    ///
732    /// When set, the `ttlMs` field is included in list responses so clients can
733    /// cache the list for up to this many milliseconds before re-fetching.
734    /// Implements SEP-2549.
735    pub fn list_ttl(mut self, ms: u64) -> Self {
736        Arc::make_mut(&mut self.inner).list_ttl_ms = Some(ms);
737        self
738    }
739
740    /// Mark the logging capability as deprecated in the server's initialize result.
741    ///
742    /// When set, the `deprecated` object is included in the `logging` capability
743    /// in the `initialize` response, signalling to clients that logging notifications
744    /// are being phased out. Implements SEP-2577.
745    pub fn logging_deprecated(mut self, info: tower_mcp_types::protocol::DeprecationInfo) -> Self {
746        Arc::make_mut(&mut self.inner).logging_deprecated = Some(info);
747        self
748    }
749
750    /// Set instructions for LLMs describing how to use this server
751    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
752        Arc::make_mut(&mut self.inner).instructions = Some(instructions.into());
753        self
754    }
755
756    /// Auto-generate instructions from registered tool, resource, and prompt descriptions.
757    ///
758    /// The instructions are generated lazily at initialization time, so this can be
759    /// called at any point in the builder chain regardless of when tools, resources,
760    /// and prompts are registered.
761    ///
762    /// If both `instructions()` and `auto_instructions()` are set, the auto-generated
763    /// instructions take precedence.
764    ///
765    /// # Example
766    ///
767    /// ```rust
768    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
769    /// use schemars::JsonSchema;
770    /// use serde::Deserialize;
771    ///
772    /// #[derive(Debug, Deserialize, JsonSchema)]
773    /// struct QueryInput { sql: String }
774    ///
775    /// let query_tool = ToolBuilder::new("query")
776    ///     .description("Execute a read-only SQL query")
777    ///     .read_only()
778    ///     .handler(|input: QueryInput| async move {
779    ///         Ok(CallToolResult::text("result"))
780    ///     })
781    ///     .build();
782    ///
783    /// let router = McpRouter::new()
784    ///     .auto_instructions()
785    ///     .tool(query_tool);
786    /// ```
787    pub fn auto_instructions(mut self) -> Self {
788        Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
789            prefix: None,
790            suffix: None,
791        });
792        self
793    }
794
795    /// Auto-generate instructions with custom prefix and/or suffix text.
796    ///
797    /// The prefix is prepended and suffix appended to the generated instructions.
798    /// See [`auto_instructions`](Self::auto_instructions) for details.
799    ///
800    /// # Example
801    ///
802    /// ```rust
803    /// use tower_mcp::McpRouter;
804    ///
805    /// let router = McpRouter::new()
806    ///     .auto_instructions_with(
807    ///         Some("This server provides database tools."),
808    ///         Some("Use 'query' for read operations and 'insert' for writes."),
809    ///     );
810    /// ```
811    pub fn auto_instructions_with(
812        mut self,
813        prefix: Option<impl Into<String>>,
814        suffix: Option<impl Into<String>>,
815    ) -> Self {
816        Arc::make_mut(&mut self.inner).auto_instructions = Some(AutoInstructionsConfig {
817            prefix: prefix.map(Into::into),
818            suffix: suffix.map(Into::into),
819        });
820        self
821    }
822
823    /// Set a human-readable title for the server
824    pub fn server_title(mut self, title: impl Into<String>) -> Self {
825        Arc::make_mut(&mut self.inner).server_title = Some(title.into());
826        self
827    }
828
829    /// Set the server description
830    pub fn server_description(mut self, description: impl Into<String>) -> Self {
831        Arc::make_mut(&mut self.inner).server_description = Some(description.into());
832        self
833    }
834
835    /// Set icons for the server
836    pub fn server_icons(mut self, icons: Vec<ToolIcon>) -> Self {
837        Arc::make_mut(&mut self.inner).server_icons = Some(icons);
838        self
839    }
840
841    /// Set the server's website URL
842    pub fn server_website_url(mut self, url: impl Into<String>) -> Self {
843        Arc::make_mut(&mut self.inner).server_website_url = Some(url.into());
844        self
845    }
846
847    /// Register a tool
848    pub fn tool(mut self, tool: Tool) -> Self {
849        Arc::make_mut(&mut self.inner)
850            .tools
851            .insert(tool.name.clone(), Arc::new(tool));
852        self
853    }
854
855    /// Conditionally register a tool.
856    ///
857    /// Registers the tool only if `condition` is `true`. This keeps fluent
858    /// builder chains intact when tools are conditionally enabled.
859    ///
860    /// # Example
861    ///
862    /// ```rust
863    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
864    /// use schemars::JsonSchema;
865    /// use serde::Deserialize;
866    ///
867    /// #[derive(Debug, Deserialize, JsonSchema)]
868    /// struct Input { value: String }
869    ///
870    /// let enable_admin = false;
871    ///
872    /// let admin_tool = ToolBuilder::new("admin")
873    ///     .description("Admin tool")
874    ///     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
875    ///     .build();
876    ///
877    /// let router = McpRouter::new()
878    ///     .tool_if(enable_admin, admin_tool);
879    /// ```
880    pub fn tool_if(self, condition: bool, tool: Tool) -> Self {
881        if condition { self.tool(tool) } else { self }
882    }
883
884    /// Register a resource
885    pub fn resource(mut self, resource: Resource) -> Self {
886        Arc::make_mut(&mut self.inner)
887            .resources
888            .insert(resource.uri.clone(), Arc::new(resource));
889        self
890    }
891
892    /// Conditionally register a resource.
893    ///
894    /// Registers the resource only if `condition` is `true`.
895    ///
896    /// # Example
897    ///
898    /// ```rust
899    /// use tower_mcp::{McpRouter, ResourceBuilder};
900    ///
901    /// let enable_config = false;
902    ///
903    /// let config = ResourceBuilder::new("config://system")
904    ///     .name("config")
905    ///     .text("secret=xxx");
906    ///
907    /// let router = McpRouter::new()
908    ///     .resource_if(enable_config, config);
909    /// ```
910    pub fn resource_if(self, condition: bool, resource: Resource) -> Self {
911        if condition {
912            self.resource(resource)
913        } else {
914            self
915        }
916    }
917
918    /// Register a resource template
919    ///
920    /// Resource templates allow dynamic resources to be matched by URI pattern.
921    /// When a client requests a resource URI that doesn't match any static
922    /// resource, the router tries to match it against registered templates.
923    ///
924    /// # Example
925    ///
926    /// ```rust
927    /// use tower_mcp::{McpRouter, ResourceTemplateBuilder};
928    /// use tower_mcp::protocol::{ReadResourceResult, ResourceContent};
929    /// use std::collections::HashMap;
930    ///
931    /// let template = ResourceTemplateBuilder::new("file:///{path}")
932    ///     .name("Project Files")
933    ///     .handler(|uri: String, vars: HashMap<String, String>| async move {
934    ///         let path = vars.get("path").unwrap_or(&String::new()).clone();
935    ///         Ok(ReadResourceResult {
936    ///             contents: vec![ResourceContent {
937    ///                 uri,
938    ///                 mime_type: Some("text/plain".to_string()),
939    ///                 text: Some(format!("Contents of {}", path)),
940    ///                 blob: None,
941    ///                 meta: None,
942    ///             }],
943    ///             meta: None,
944    ///         })
945    ///     });
946    ///
947    /// let router = McpRouter::new()
948    ///     .resource_template(template);
949    /// ```
950    pub fn resource_template(mut self, template: ResourceTemplate) -> Self {
951        Arc::make_mut(&mut self.inner)
952            .resource_templates
953            .push(Arc::new(template));
954        self
955    }
956
957    /// Register a prompt
958    pub fn prompt(mut self, prompt: Prompt) -> Self {
959        Arc::make_mut(&mut self.inner)
960            .prompts
961            .insert(prompt.name.clone(), Arc::new(prompt));
962        self
963    }
964
965    /// Conditionally register a prompt.
966    ///
967    /// Registers the prompt only if `condition` is `true`.
968    ///
969    /// # Example
970    ///
971    /// ```rust
972    /// use tower_mcp::{McpRouter, PromptBuilder};
973    ///
974    /// let enable_debug = false;
975    ///
976    /// let debug_prompt = PromptBuilder::new("debug")
977    ///     .description("Debug prompt")
978    ///     .user_message("Debug mode enabled");
979    ///
980    /// let router = McpRouter::new()
981    ///     .prompt_if(enable_debug, debug_prompt);
982    /// ```
983    pub fn prompt_if(self, condition: bool, prompt: Prompt) -> Self {
984        if condition { self.prompt(prompt) } else { self }
985    }
986
987    /// Register multiple tools at once.
988    ///
989    /// # Example
990    ///
991    /// ```rust
992    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
993    /// use schemars::JsonSchema;
994    /// use serde::Deserialize;
995    ///
996    /// #[derive(Debug, Deserialize, JsonSchema)]
997    /// struct Input { value: String }
998    ///
999    /// let tools = vec![
1000    ///     ToolBuilder::new("a")
1001    ///         .description("Tool A")
1002    ///         .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1003    ///         .build(),
1004    ///     ToolBuilder::new("b")
1005    ///         .description("Tool B")
1006    ///         .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1007    ///         .build(),
1008    /// ];
1009    ///
1010    /// let router = McpRouter::new().tools(tools);
1011    /// ```
1012    pub fn tools(self, tools: impl IntoIterator<Item = Tool>) -> Self {
1013        tools
1014            .into_iter()
1015            .fold(self, |router, tool| router.tool(tool))
1016    }
1017
1018    /// Conditionally register multiple tools at once.
1019    ///
1020    /// Registers all tools only if `condition` is `true`.
1021    pub fn tools_if(self, condition: bool, tools: impl IntoIterator<Item = Tool>) -> Self {
1022        if condition { self.tools(tools) } else { self }
1023    }
1024
1025    /// Register multiple resources at once.
1026    ///
1027    /// # Example
1028    ///
1029    /// ```rust
1030    /// use tower_mcp::{McpRouter, ResourceBuilder};
1031    ///
1032    /// let resources = vec![
1033    ///     ResourceBuilder::new("file:///a.txt")
1034    ///         .name("File A")
1035    ///         .text("contents a"),
1036    ///     ResourceBuilder::new("file:///b.txt")
1037    ///         .name("File B")
1038    ///         .text("contents b"),
1039    /// ];
1040    ///
1041    /// let router = McpRouter::new().resources(resources);
1042    /// ```
1043    pub fn resources(self, resources: impl IntoIterator<Item = Resource>) -> Self {
1044        resources
1045            .into_iter()
1046            .fold(self, |router, resource| router.resource(resource))
1047    }
1048
1049    /// Conditionally register multiple resources at once.
1050    ///
1051    /// Registers all resources only if `condition` is `true`.
1052    pub fn resources_if(
1053        self,
1054        condition: bool,
1055        resources: impl IntoIterator<Item = Resource>,
1056    ) -> Self {
1057        if condition {
1058            self.resources(resources)
1059        } else {
1060            self
1061        }
1062    }
1063
1064    /// Register multiple prompts at once.
1065    ///
1066    /// # Example
1067    ///
1068    /// ```rust
1069    /// use tower_mcp::{McpRouter, PromptBuilder};
1070    ///
1071    /// let prompts = vec![
1072    ///     PromptBuilder::new("greet")
1073    ///         .description("Greet someone")
1074    ///         .user_message("Hello!"),
1075    ///     PromptBuilder::new("farewell")
1076    ///         .description("Say goodbye")
1077    ///         .user_message("Goodbye!"),
1078    /// ];
1079    ///
1080    /// let router = McpRouter::new().prompts(prompts);
1081    /// ```
1082    pub fn prompts(self, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1083        prompts
1084            .into_iter()
1085            .fold(self, |router, prompt| router.prompt(prompt))
1086    }
1087
1088    /// Conditionally register multiple prompts at once.
1089    ///
1090    /// Registers all prompts only if `condition` is `true`.
1091    pub fn prompts_if(self, condition: bool, prompts: impl IntoIterator<Item = Prompt>) -> Self {
1092        if condition {
1093            self.prompts(prompts)
1094        } else {
1095            self
1096        }
1097    }
1098
1099    /// Merge another router's capabilities into this one.
1100    ///
1101    /// This combines all tools, resources, resource templates, and prompts from
1102    /// the other router into this router. Uses "last wins" semantics for conflicts,
1103    /// meaning if both routers have a tool/resource/prompt with the same name,
1104    /// the one from `other` will replace the one in `self`.
1105    ///
1106    /// Server info, instructions, filters, and other router-level configuration
1107    /// are NOT merged - only the root router's settings are used.
1108    ///
1109    /// # Example
1110    ///
1111    /// ```rust
1112    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult, ResourceBuilder};
1113    /// use schemars::JsonSchema;
1114    /// use serde::Deserialize;
1115    ///
1116    /// #[derive(Debug, Deserialize, JsonSchema)]
1117    /// struct Input { value: String }
1118    ///
1119    /// // Create a router with database tools
1120    /// let db_tools = McpRouter::new()
1121    ///     .tool(
1122    ///         ToolBuilder::new("query")
1123    ///             .description("Query the database")
1124    ///             .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1125    ///             .build()
1126    ///     );
1127    ///
1128    /// // Create a router with API tools
1129    /// let api_tools = McpRouter::new()
1130    ///     .tool(
1131    ///         ToolBuilder::new("fetch")
1132    ///             .description("Fetch from API")
1133    ///             .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1134    ///             .build()
1135    ///     );
1136    ///
1137    /// // Merge them together
1138    /// let router = McpRouter::new()
1139    ///     .server_info("combined", "1.0")
1140    ///     .merge(db_tools)
1141    ///     .merge(api_tools);
1142    /// ```
1143    pub fn merge(mut self, other: McpRouter) -> Self {
1144        let inner = Arc::make_mut(&mut self.inner);
1145        let other_inner = other.inner;
1146
1147        // Merge tools (last wins)
1148        for (name, tool) in &other_inner.tools {
1149            inner.tools.insert(name.clone(), tool.clone());
1150        }
1151
1152        // Merge resources (last wins)
1153        for (uri, resource) in &other_inner.resources {
1154            inner.resources.insert(uri.clone(), resource.clone());
1155        }
1156
1157        // Merge resource templates (append - no deduplication since templates
1158        // can have complex matching behavior)
1159        for template in &other_inner.resource_templates {
1160            inner.resource_templates.push(template.clone());
1161        }
1162
1163        // Merge prompts (last wins)
1164        for (name, prompt) in &other_inner.prompts {
1165            inner.prompts.insert(name.clone(), prompt.clone());
1166        }
1167
1168        self
1169    }
1170
1171    /// Nest another router's capabilities under a prefix.
1172    ///
1173    /// This is similar to `merge()`, but all tool names from the nested router
1174    /// are prefixed with the given string and a dot separator. For example,
1175    /// nesting with prefix "db" will turn a tool named "query" into "db.query".
1176    ///
1177    /// Resources, resource templates, and prompts are merged without modification
1178    /// since they use URIs rather than simple names for identification.
1179    ///
1180    /// # Example
1181    ///
1182    /// ```rust
1183    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
1184    /// use schemars::JsonSchema;
1185    /// use serde::Deserialize;
1186    ///
1187    /// #[derive(Debug, Deserialize, JsonSchema)]
1188    /// struct Input { value: String }
1189    ///
1190    /// // Create a router with database tools
1191    /// let db_tools = McpRouter::new()
1192    ///     .tool(
1193    ///         ToolBuilder::new("query")
1194    ///             .description("Query the database")
1195    ///             .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1196    ///             .build()
1197    ///     )
1198    ///     .tool(
1199    ///         ToolBuilder::new("insert")
1200    ///             .description("Insert into database")
1201    ///             .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1202    ///             .build()
1203    ///     );
1204    ///
1205    /// // Nest under "db" prefix - tools become "db.query" and "db.insert"
1206    /// let router = McpRouter::new()
1207    ///     .server_info("combined", "1.0")
1208    ///     .nest("db", db_tools);
1209    /// ```
1210    pub fn nest(mut self, prefix: impl Into<String>, other: McpRouter) -> Self {
1211        let prefix = prefix.into();
1212        let inner = Arc::make_mut(&mut self.inner);
1213        let other_inner = other.inner;
1214
1215        // Nest tools with prefix
1216        for tool in other_inner.tools.values() {
1217            let prefixed_tool = tool.with_name_prefix(&prefix);
1218            inner
1219                .tools
1220                .insert(prefixed_tool.name.clone(), Arc::new(prefixed_tool));
1221        }
1222
1223        // Merge resources (no prefix - URIs are already namespaced)
1224        for (uri, resource) in &other_inner.resources {
1225            inner.resources.insert(uri.clone(), resource.clone());
1226        }
1227
1228        // Merge resource templates (no prefix)
1229        for template in &other_inner.resource_templates {
1230            inner.resource_templates.push(template.clone());
1231        }
1232
1233        // Merge prompts (no prefix - could be added in future if needed)
1234        for (name, prompt) in &other_inner.prompts {
1235            inner.prompts.insert(name.clone(), prompt.clone());
1236        }
1237
1238        self
1239    }
1240
1241    /// Register a completion handler for `completion/complete` requests.
1242    ///
1243    /// The handler receives `CompleteParams` containing the reference (prompt or resource)
1244    /// and the argument being completed, and should return completion suggestions.
1245    ///
1246    /// # Example
1247    ///
1248    /// ```rust
1249    /// use tower_mcp::{McpRouter, CompleteResult};
1250    /// use tower_mcp::protocol::{CompleteParams, CompletionReference};
1251    ///
1252    /// let router = McpRouter::new()
1253    ///     .completion_handler(|params: CompleteParams| async move {
1254    ///         // Provide completions based on the reference and argument
1255    ///         match params.reference {
1256    ///             CompletionReference::Prompt { name } => {
1257    ///                 // Return prompt argument completions
1258    ///                 Ok(CompleteResult::new(vec!["option1".to_string(), "option2".to_string()]))
1259    ///             }
1260    ///             CompletionReference::Resource { uri } => {
1261    ///                 // Return resource URI completions
1262    ///                 Ok(CompleteResult::new(vec![]))
1263    ///             }
1264    ///             _ => Ok(CompleteResult::new(vec![])),
1265    ///         }
1266    ///     });
1267    /// ```
1268    pub fn completion_handler<F, Fut>(mut self, handler: F) -> Self
1269    where
1270        F: Fn(CompleteParams) -> Fut + Send + Sync + 'static,
1271        Fut: Future<Output = Result<CompleteResult>> + Send + 'static,
1272    {
1273        Arc::make_mut(&mut self.inner).completion_handler =
1274            Some(Arc::new(move |params| Box::pin(handler(params))));
1275        self
1276    }
1277
1278    /// Set a filter for tools based on session state.
1279    ///
1280    /// The filter determines which tools are visible to each session. Tools that
1281    /// don't pass the filter will not appear in `tools/list` responses and will
1282    /// return an error if called directly.
1283    ///
1284    /// # Example
1285    ///
1286    /// ```rust
1287    /// use tower_mcp::{McpRouter, ToolBuilder, CallToolResult, CapabilityFilter, Tool, Filterable};
1288    /// use schemars::JsonSchema;
1289    /// use serde::Deserialize;
1290    ///
1291    /// #[derive(Debug, Deserialize, JsonSchema)]
1292    /// struct Input { value: String }
1293    ///
1294    /// let public_tool = ToolBuilder::new("public")
1295    ///     .description("Available to everyone")
1296    ///     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1297    ///     .build();
1298    ///
1299    /// let admin_tool = ToolBuilder::new("admin")
1300    ///     .description("Admin only")
1301    ///     .handler(|i: Input| async move { Ok(CallToolResult::text(&i.value)) })
1302    ///     .build();
1303    ///
1304    /// let router = McpRouter::new()
1305    ///     .tool(public_tool)
1306    ///     .tool(admin_tool)
1307    ///     .tool_filter(CapabilityFilter::new(|_session, tool: &Tool| {
1308    ///         // In real code, check session.extensions() for auth claims
1309    ///         tool.name() != "admin"
1310    ///     }));
1311    /// ```
1312    pub fn tool_filter(mut self, filter: ToolFilter) -> Self {
1313        Arc::make_mut(&mut self.inner).tool_filter = Some(filter);
1314        self
1315    }
1316
1317    /// Set a filter for resources based on session state.
1318    ///
1319    /// The filter receives the current session state and each resource, returning
1320    /// `true` if the resource should be visible to this session. Resources that
1321    /// don't pass the filter will not appear in `resources/list` responses and will
1322    /// return an error if read directly.
1323    ///
1324    /// # Example
1325    ///
1326    /// ```rust
1327    /// use tower_mcp::{McpRouter, ResourceBuilder, ReadResourceResult, CapabilityFilter, Resource, Filterable};
1328    ///
1329    /// let public_resource = ResourceBuilder::new("file:///public.txt")
1330    ///     .name("Public File")
1331    ///     .description("Available to everyone")
1332    ///     .text("public content");
1333    ///
1334    /// let secret_resource = ResourceBuilder::new("file:///secret.txt")
1335    ///     .name("Secret File")
1336    ///     .description("Admin only")
1337    ///     .text("secret content");
1338    ///
1339    /// let router = McpRouter::new()
1340    ///     .resource(public_resource)
1341    ///     .resource(secret_resource)
1342    ///     .resource_filter(CapabilityFilter::new(|_session, resource: &Resource| {
1343    ///         // In real code, check session.extensions() for auth claims
1344    ///         !resource.name().contains("Secret")
1345    ///     }));
1346    /// ```
1347    pub fn resource_filter(mut self, filter: ResourceFilter) -> Self {
1348        Arc::make_mut(&mut self.inner).resource_filter = Some(filter);
1349        self
1350    }
1351
1352    /// Set a filter for prompts based on session state.
1353    ///
1354    /// The filter receives the current session state and each prompt, returning
1355    /// `true` if the prompt should be visible to this session. Prompts that
1356    /// don't pass the filter will not appear in `prompts/list` responses and will
1357    /// return an error if accessed directly.
1358    ///
1359    /// # Example
1360    ///
1361    /// ```rust
1362    /// use tower_mcp::{McpRouter, PromptBuilder, CapabilityFilter, Prompt, Filterable};
1363    ///
1364    /// let public_prompt = PromptBuilder::new("greeting")
1365    ///     .description("A friendly greeting")
1366    ///     .user_message("Hello!");
1367    ///
1368    /// let admin_prompt = PromptBuilder::new("system_debug")
1369    ///     .description("Admin debugging prompt")
1370    ///     .user_message("Debug info");
1371    ///
1372    /// let router = McpRouter::new()
1373    ///     .prompt(public_prompt)
1374    ///     .prompt(admin_prompt)
1375    ///     .prompt_filter(CapabilityFilter::new(|_session, prompt: &Prompt| {
1376    ///         // In real code, check session.extensions() for auth claims
1377    ///         !prompt.name().contains("system")
1378    ///     }));
1379    /// ```
1380    pub fn prompt_filter(mut self, filter: PromptFilter) -> Self {
1381        Arc::make_mut(&mut self.inner).prompt_filter = Some(filter);
1382        self
1383    }
1384
1385    /// Get access to the session state
1386    pub fn session(&self) -> &SessionState {
1387        &self.session
1388    }
1389
1390    /// Send a log message notification to the client
1391    ///
1392    /// This sends a `notifications/message` notification with the given parameters.
1393    /// Returns `true` if the notification was sent, `false` if no notification channel
1394    /// is configured.
1395    ///
1396    /// # Example
1397    ///
1398    /// ```rust,ignore
1399    /// use tower_mcp::protocol::{LogLevel, LoggingMessageParams};
1400    ///
1401    /// // Simple info message
1402    /// router.log(LoggingMessageParams::new(LogLevel::Info,
1403    ///     serde_json::json!({"message": "Operation completed"})
1404    /// ));
1405    ///
1406    /// // Error with logger name
1407    /// router.log(LoggingMessageParams::new(LogLevel::Error,
1408    ///     serde_json::json!({"error": "Connection failed"}))
1409    ///     .with_logger("database"));
1410    /// ```
1411    pub fn log(&self, params: LoggingMessageParams) -> bool {
1412        let Some(tx) = &self.inner.notification_tx else {
1413            return false;
1414        };
1415        tx.try_send(ServerNotification::LogMessage(params)).is_ok()
1416    }
1417
1418    /// Send an info-level log message
1419    ///
1420    /// Convenience method for sending an info log with a message string.
1421    pub fn log_info(&self, message: &str) -> bool {
1422        self.log(LoggingMessageParams::new(
1423            LogLevel::Info,
1424            serde_json::json!({ "message": message }),
1425        ))
1426    }
1427
1428    /// Send a warning-level log message
1429    pub fn log_warning(&self, message: &str) -> bool {
1430        self.log(LoggingMessageParams::new(
1431            LogLevel::Warning,
1432            serde_json::json!({ "message": message }),
1433        ))
1434    }
1435
1436    /// Send an error-level log message
1437    pub fn log_error(&self, message: &str) -> bool {
1438        self.log(LoggingMessageParams::new(
1439            LogLevel::Error,
1440            serde_json::json!({ "message": message }),
1441        ))
1442    }
1443
1444    /// Send a debug-level log message
1445    pub fn log_debug(&self, message: &str) -> bool {
1446        self.log(LoggingMessageParams::new(
1447            LogLevel::Debug,
1448            serde_json::json!({ "message": message }),
1449        ))
1450    }
1451
1452    /// Check if a resource URI is currently subscribed
1453    pub fn is_subscribed(&self, uri: &str) -> bool {
1454        if let Ok(subs) = self.inner.subscriptions.read() {
1455            return subs.contains(uri);
1456        }
1457        false
1458    }
1459
1460    /// Get a list of all subscribed resource URIs
1461    pub fn subscribed_uris(&self) -> Vec<String> {
1462        if let Ok(subs) = self.inner.subscriptions.read() {
1463            return subs.iter().cloned().collect();
1464        }
1465        Vec::new()
1466    }
1467
1468    /// Subscribe to a resource URI
1469    fn subscribe(&self, uri: &str) -> bool {
1470        if let Ok(mut subs) = self.inner.subscriptions.write() {
1471            return subs.insert(uri.to_string());
1472        }
1473        false
1474    }
1475
1476    /// Unsubscribe from a resource URI
1477    fn unsubscribe(&self, uri: &str) -> bool {
1478        if let Ok(mut subs) = self.inner.subscriptions.write() {
1479            return subs.remove(uri);
1480        }
1481        false
1482    }
1483
1484    /// Notify clients that a subscribed resource has been updated
1485    ///
1486    /// Only sends the notification if the resource is currently subscribed.
1487    /// Returns `true` if the notification was sent.
1488    pub fn notify_resource_updated(&self, uri: &str) -> bool {
1489        // Only notify if the resource is subscribed
1490        if !self.is_subscribed(uri) {
1491            return false;
1492        }
1493
1494        let Some(tx) = &self.inner.notification_tx else {
1495            return false;
1496        };
1497        tx.try_send(ServerNotification::ResourceUpdated {
1498            uri: uri.to_string(),
1499        })
1500        .is_ok()
1501    }
1502
1503    /// Notify clients that the list of available resources has changed
1504    ///
1505    /// Returns `true` if the notification was sent.
1506    pub fn notify_resources_list_changed(&self) -> bool {
1507        let Some(tx) = &self.inner.notification_tx else {
1508            return false;
1509        };
1510        tx.try_send(ServerNotification::ResourcesListChanged)
1511            .is_ok()
1512    }
1513
1514    /// Notify clients that the list of available tools has changed
1515    ///
1516    /// Returns `true` if the notification was sent.
1517    pub fn notify_tools_list_changed(&self) -> bool {
1518        let Some(tx) = &self.inner.notification_tx else {
1519            return false;
1520        };
1521        tx.try_send(ServerNotification::ToolsListChanged).is_ok()
1522    }
1523
1524    /// Notify clients that the list of available prompts has changed
1525    ///
1526    /// Returns `true` if the notification was sent.
1527    pub fn notify_prompts_list_changed(&self) -> bool {
1528        let Some(tx) = &self.inner.notification_tx else {
1529            return false;
1530        };
1531        tx.try_send(ServerNotification::PromptsListChanged).is_ok()
1532    }
1533
1534    /// Disable a tool by name. Disabled tools are hidden from `tools/list`
1535    /// and return a method-not-found error from `tools/call`, but the tool
1536    /// definition stays attached to the router and can be flipped back on
1537    /// with [`enable_tool`](Self::enable_tool).
1538    ///
1539    /// State is shared across all clones produced by
1540    /// [`with_fresh_session`](Self::with_fresh_session), so flipping it once
1541    /// affects every connected session at the next request boundary. Call
1542    /// [`notify_tools_list_changed`](Self::notify_tools_list_changed) to nudge
1543    /// clients to re-fetch.
1544    pub fn disable_tool(&self, name: impl Into<String>) {
1545        let mut set = self.inner.disabled_tools.write().unwrap();
1546        set.insert(name.into());
1547    }
1548
1549    /// Re-enable a previously disabled tool. No-op if the tool was not
1550    /// disabled.
1551    pub fn enable_tool(&self, name: &str) {
1552        let mut set = self.inner.disabled_tools.write().unwrap();
1553        set.remove(name);
1554    }
1555
1556    /// Returns `true` if the named tool is currently enabled (i.e. not in
1557    /// the disabled set). Returns `true` even for unknown tool names; this
1558    /// only reports disable state, not registration.
1559    pub fn is_tool_enabled(&self, name: &str) -> bool {
1560        !self.inner.disabled_tools.read().unwrap().contains(name)
1561    }
1562
1563    /// Disable a resource by URI. Disabled resources are hidden from
1564    /// `resources/list` and return a not-found error from `resources/read`.
1565    pub fn disable_resource(&self, uri: impl Into<String>) {
1566        let mut set = self.inner.disabled_resources.write().unwrap();
1567        set.insert(uri.into());
1568    }
1569
1570    /// Re-enable a previously disabled resource.
1571    pub fn enable_resource(&self, uri: &str) {
1572        let mut set = self.inner.disabled_resources.write().unwrap();
1573        set.remove(uri);
1574    }
1575
1576    /// Returns `true` if the resource at this URI is currently enabled.
1577    pub fn is_resource_enabled(&self, uri: &str) -> bool {
1578        !self.inner.disabled_resources.read().unwrap().contains(uri)
1579    }
1580
1581    /// Disable a prompt by name. Disabled prompts are hidden from
1582    /// `prompts/list` and return a method-not-found error from `prompts/get`.
1583    pub fn disable_prompt(&self, name: impl Into<String>) {
1584        let mut set = self.inner.disabled_prompts.write().unwrap();
1585        set.insert(name.into());
1586    }
1587
1588    /// Re-enable a previously disabled prompt.
1589    pub fn enable_prompt(&self, name: &str) {
1590        let mut set = self.inner.disabled_prompts.write().unwrap();
1591        set.remove(name);
1592    }
1593
1594    /// Returns `true` if the named prompt is currently enabled.
1595    pub fn is_prompt_enabled(&self, name: &str) -> bool {
1596        !self.inner.disabled_prompts.read().unwrap().contains(name)
1597    }
1598
1599    /// Get server capabilities based on registered handlers
1600    fn capabilities(&self) -> ServerCapabilities {
1601        let has_resources =
1602            !self.inner.resources.is_empty() || !self.inner.resource_templates.is_empty();
1603        let has_notifications = self.inner.notification_tx.is_some();
1604
1605        #[cfg(feature = "dynamic-tools")]
1606        let has_dynamic_tools = self.inner.dynamic_tools.is_some();
1607        #[cfg(not(feature = "dynamic-tools"))]
1608        let has_dynamic_tools = false;
1609
1610        #[cfg(feature = "dynamic-tools")]
1611        let has_dynamic_prompts = self.inner.dynamic_prompts.is_some();
1612        #[cfg(not(feature = "dynamic-tools"))]
1613        let has_dynamic_prompts = false;
1614
1615        #[cfg(feature = "dynamic-tools")]
1616        let has_dynamic_resources = self.inner.dynamic_resources.is_some()
1617            || self.inner.dynamic_resource_templates.is_some();
1618        #[cfg(not(feature = "dynamic-tools"))]
1619        let has_dynamic_resources = false;
1620
1621        ServerCapabilities {
1622            tools: if self.inner.tools.is_empty() && !has_dynamic_tools {
1623                None
1624            } else {
1625                Some(ToolsCapability {
1626                    list_changed: has_notifications,
1627                })
1628            },
1629            resources: if has_resources || has_dynamic_resources {
1630                Some(ResourcesCapability {
1631                    subscribe: true,
1632                    list_changed: has_notifications,
1633                })
1634            } else {
1635                None
1636            },
1637            prompts: if self.inner.prompts.is_empty() && !has_dynamic_prompts {
1638                None
1639            } else {
1640                Some(PromptsCapability {
1641                    list_changed: has_notifications,
1642                })
1643            },
1644            // Always advertise logging capability when notification channel is configured
1645            logging: if self.inner.notification_tx.is_some() {
1646                Some(LoggingCapability {
1647                    deprecated: self.inner.logging_deprecated.clone(),
1648                })
1649            } else {
1650                None
1651            },
1652            // Tasks capability is advertised if any tool supports tasks.
1653            // SEP-2663 moves the declaration to `capabilities.extensions`
1654            // under the reverse-DNS key `io.modelcontextprotocol/tasks`; we
1655            // continue to set the legacy top-level `tasks` field for back-compat
1656            // with 2025-11-25 clients that key off it.
1657            tasks: {
1658                let has_task_support = self
1659                    .inner
1660                    .tools
1661                    .values()
1662                    .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1663                if has_task_support {
1664                    Some(TasksCapability {
1665                        list: Some(TasksListCapability {}),
1666                        cancel: Some(TasksCancelCapability {}),
1667                        requests: Some(TasksRequestsCapability {
1668                            tools: Some(TasksToolsRequestsCapability {
1669                                call: Some(TasksToolsCallCapability {}),
1670                            }),
1671                        }),
1672                    })
1673                } else {
1674                    None
1675                }
1676            },
1677            // Completions capability when a handler is registered
1678            completions: if self.inner.completion_handler.is_some() {
1679                Some(CompletionsCapability::default())
1680            } else {
1681                None
1682            },
1683            experimental: None,
1684            extensions: {
1685                let has_task_support = self
1686                    .inner
1687                    .tools
1688                    .values()
1689                    .any(|t| !matches!(t.task_support, TaskSupportMode::Forbidden));
1690                if has_task_support {
1691                    let mut map = std::collections::HashMap::new();
1692                    map.insert(
1693                        tower_mcp_types::protocol::TASKS_EXTENSION_ID.to_string(),
1694                        serde_json::json!({}),
1695                    );
1696                    Some(map)
1697                } else {
1698                    None
1699                }
1700            },
1701        }
1702    }
1703
1704    /// Handle an MCP request
1705    async fn handle(
1706        &self,
1707        request_id: RequestId,
1708        request: McpRequest,
1709        extensions: Extensions,
1710    ) -> Result<McpResponse> {
1711        // Enforce session state - reject requests before initialization
1712        let method = request.method_name();
1713        if !self.session.is_request_allowed(method) {
1714            tracing::warn!(
1715                method = %method,
1716                phase = ?self.session.phase(),
1717                "Request rejected: session not initialized"
1718            );
1719            return Err(Error::JsonRpc(JsonRpcError::invalid_request(format!(
1720                "Session not initialized. Only 'initialize' and 'ping' are allowed before initialization. Got: {}",
1721                method
1722            ))));
1723        }
1724
1725        match request {
1726            McpRequest::Initialize(params) => {
1727                tracing::info!(
1728                    client = %params.client_info.name,
1729                    version = %params.client_info.version,
1730                    "Client initializing"
1731                );
1732
1733                // Protocol version negotiation: respond with same version if supported,
1734                // otherwise respond with our latest supported version
1735                let protocol_version = if crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1736                    .contains(&params.protocol_version.as_str())
1737                {
1738                    params.protocol_version
1739                } else {
1740                    crate::protocol::LATEST_PROTOCOL_VERSION.to_string()
1741                };
1742
1743                // Transition session state to Initializing
1744                self.session.mark_initializing();
1745
1746                Ok(McpResponse::Initialize(InitializeResult {
1747                    protocol_version,
1748                    capabilities: self.capabilities(),
1749                    server_info: Implementation {
1750                        name: self.inner.server_name.clone(),
1751                        version: self.inner.server_version.clone(),
1752                        title: self.inner.server_title.clone(),
1753                        description: self.inner.server_description.clone(),
1754                        icons: self.inner.server_icons.clone(),
1755                        website_url: self.inner.server_website_url.clone(),
1756                        meta: None,
1757                    },
1758                    instructions: if let Some(config) = &self.inner.auto_instructions {
1759                        Some(self.inner.generate_instructions(config))
1760                    } else {
1761                        self.inner.instructions.clone()
1762                    },
1763                    meta: None,
1764                }))
1765            }
1766
1767            McpRequest::Discover(_) => {
1768                // SEP-2575 server/discover -- stateless capability advertisement.
1769                // Unlike initialize, this does NOT transition session state and
1770                // does not require a session at all. Returns the same capability
1771                // surface plus the full set of protocol versions we can speak,
1772                // so clients can pick one and signal it via MCP-Protocol-Version
1773                // on subsequent requests.
1774                tracing::debug!("Stateless server/discover request");
1775                Ok(McpResponse::Discover(DiscoverResult {
1776                    supported_versions: crate::protocol::SUPPORTED_PROTOCOL_VERSIONS
1777                        .iter()
1778                        .map(|v| (*v).to_string())
1779                        .collect(),
1780                    capabilities: self.capabilities(),
1781                    server_info: Implementation {
1782                        name: self.inner.server_name.clone(),
1783                        version: self.inner.server_version.clone(),
1784                        title: self.inner.server_title.clone(),
1785                        description: self.inner.server_description.clone(),
1786                        icons: self.inner.server_icons.clone(),
1787                        website_url: self.inner.server_website_url.clone(),
1788                        meta: None,
1789                    },
1790                    instructions: if let Some(config) = &self.inner.auto_instructions {
1791                        Some(self.inner.generate_instructions(config))
1792                    } else {
1793                        self.inner.instructions.clone()
1794                    },
1795                    meta: None,
1796                }))
1797            }
1798
1799            McpRequest::ListTools(params) => {
1800                let filter = self.inner.tool_filter.as_ref();
1801                let disabled = self.inner.disabled_tools.read().unwrap().clone();
1802                let is_visible = |t: &Tool| {
1803                    !disabled.contains(&t.name)
1804                        && filter
1805                            .map(|f| f.is_visible(&self.session, t))
1806                            .unwrap_or(true)
1807                };
1808
1809                // Collect static tools
1810                let mut tools: Vec<ToolDefinition> = self
1811                    .inner
1812                    .tools
1813                    .values()
1814                    .filter(|t| is_visible(t))
1815                    .map(|t| t.definition())
1816                    .collect();
1817
1818                // Merge dynamic tools (static tools win on name collision)
1819                #[cfg(feature = "dynamic-tools")]
1820                if let Some(ref dynamic) = self.inner.dynamic_tools {
1821                    let static_names: HashSet<String> =
1822                        tools.iter().map(|t| t.name.clone()).collect();
1823                    for t in dynamic.list() {
1824                        if !static_names.contains(&t.name) && is_visible(&t) {
1825                            tools.push(t.definition());
1826                        }
1827                    }
1828                }
1829
1830                tools.sort_by(|a, b| a.name.cmp(&b.name));
1831
1832                let (tools, next_cursor) =
1833                    paginate(tools, params.cursor.as_deref(), self.inner.page_size)?;
1834
1835                Ok(McpResponse::ListTools(ListToolsResult {
1836                    tools,
1837                    next_cursor,
1838                    ttl_ms: self.inner.list_ttl_ms,
1839                    cache_scope: None,
1840                    meta: None,
1841                }))
1842            }
1843
1844            McpRequest::CallTool(params) => {
1845                // Disabled tools are reported as if they don't exist.
1846                if self
1847                    .inner
1848                    .disabled_tools
1849                    .read()
1850                    .unwrap()
1851                    .contains(&params.name)
1852                {
1853                    tracing::info!(
1854                        target: "mcp::tools",
1855                        tool = %params.name,
1856                        status = "disabled",
1857                        "tool call completed"
1858                    );
1859                    return Err(Error::JsonRpc(JsonRpcError::method_not_found(&params.name)));
1860                }
1861
1862                // Look up static tools first, then dynamic
1863                let tool = self.inner.tools.get(&params.name).cloned();
1864                #[cfg(feature = "dynamic-tools")]
1865                let tool = tool.or_else(|| {
1866                    self.inner
1867                        .dynamic_tools
1868                        .as_ref()
1869                        .and_then(|d| d.get(&params.name))
1870                });
1871
1872                let tool = match tool {
1873                    Some(t) => t,
1874                    None => {
1875                        tracing::info!(
1876                            target: "mcp::tools",
1877                            tool = %params.name,
1878                            status = "not_found",
1879                            "tool call completed"
1880                        );
1881                        return Err(Error::JsonRpc(JsonRpcError::method_not_found(&params.name)));
1882                    }
1883                };
1884
1885                // Check tool filter if configured
1886                if let Some(filter) = &self.inner.tool_filter
1887                    && !filter.is_visible(&self.session, &tool)
1888                {
1889                    tracing::info!(
1890                        target: "mcp::tools",
1891                        tool = %params.name,
1892                        status = "denied",
1893                        "tool call completed"
1894                    );
1895                    return Err(filter.denial_error(&params.name));
1896                }
1897
1898                if let Some(task_params) = params.task {
1899                    // Task-augmented request: validate task_support != Forbidden
1900                    if matches!(tool.task_support, TaskSupportMode::Forbidden) {
1901                        return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1902                            "Tool '{}' does not support async tasks",
1903                            params.name
1904                        ))));
1905                    }
1906
1907                    // Create the task
1908                    let (task_id, cancellation_token) = self.inner.task_store.create_task(
1909                        &params.name,
1910                        params.arguments.clone(),
1911                        task_params.ttl,
1912                    );
1913
1914                    tracing::info!(task_id = %task_id, tool = %params.name, "Created async task");
1915
1916                    // Create a context for the async task execution
1917                    let progress_token = params.meta.and_then(|m| m.progress_token);
1918                    let ctx = self.create_context_with_extensions(
1919                        request_id,
1920                        progress_token,
1921                        &extensions,
1922                    );
1923
1924                    // Spawn the task execution in the background
1925                    let task_store = self.inner.task_store.clone();
1926                    let tool = tool.clone();
1927                    let arguments = params.arguments;
1928                    let task_id_clone = task_id.clone();
1929
1930                    let tool_name = params.name.clone();
1931                    tokio::spawn(async move {
1932                        // Check for cancellation before starting
1933                        if cancellation_token.is_cancelled() {
1934                            tracing::debug!(task_id = %task_id_clone, "Task cancelled before execution");
1935                            return;
1936                        }
1937
1938                        // Execute the tool
1939                        let start = std::time::Instant::now();
1940                        let result = tool.call_with_context(ctx, arguments).await;
1941                        let duration_ms = start.elapsed().as_secs_f64() * 1000.0;
1942
1943                        if cancellation_token.is_cancelled() {
1944                            tracing::debug!(task_id = %task_id_clone, "Task cancelled during execution");
1945                        } else if result.is_error {
1946                            // Tool returned an error result
1947                            let error_msg = result.first_text().unwrap_or("Tool execution failed");
1948                            task_store.fail_task(&task_id_clone, error_msg);
1949                            tracing::info!(
1950                                target: "mcp::tools",
1951                                tool = %tool_name,
1952                                task_id = %task_id_clone,
1953                                duration_ms,
1954                                status = "error",
1955                                error = %error_msg,
1956                                "tool call completed"
1957                            );
1958                        } else {
1959                            task_store.complete_task(&task_id_clone, result);
1960                            tracing::info!(
1961                                target: "mcp::tools",
1962                                tool = %tool_name,
1963                                task_id = %task_id_clone,
1964                                duration_ms,
1965                                status = "success",
1966                                "tool call completed"
1967                            );
1968                        }
1969                    });
1970
1971                    let task = self.inner.task_store.get_task(&task_id).ok_or_else(|| {
1972                        Error::JsonRpc(JsonRpcError::internal_error(
1973                            "Failed to retrieve created task",
1974                        ))
1975                    })?;
1976
1977                    Ok(McpResponse::CreateTask(CreateTaskResult::new(task)))
1978                } else {
1979                    // Synchronous request: validate task_support != Required
1980                    if matches!(tool.task_support, TaskSupportMode::Required) {
1981                        return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1982                            "Tool '{}' requires async task execution (include 'task' in params)",
1983                            params.name
1984                        ))));
1985                    }
1986
1987                    // Extract progress token from request metadata
1988                    let progress_token = params.meta.and_then(|m| m.progress_token);
1989                    let ctx = self.create_context_with_extensions(
1990                        request_id,
1991                        progress_token,
1992                        &extensions,
1993                    );
1994
1995                    let start = std::time::Instant::now();
1996                    let result = tool.call_with_context(ctx, params.arguments).await;
1997                    let duration_ms = start.elapsed().as_secs_f64() * 1000.0;
1998
1999                    if result.is_error {
2000                        tracing::info!(
2001                            target: "mcp::tools",
2002                            tool = %params.name,
2003                            duration_ms,
2004                            status = "error",
2005                            "tool call completed"
2006                        );
2007                    } else {
2008                        tracing::info!(
2009                            target: "mcp::tools",
2010                            tool = %params.name,
2011                            duration_ms,
2012                            status = "success",
2013                            "tool call completed"
2014                        );
2015                    }
2016
2017                    Ok(McpResponse::CallTool(result))
2018                }
2019            }
2020
2021            McpRequest::ListResources(params) => {
2022                let disabled = self.inner.disabled_resources.read().unwrap().clone();
2023                let is_visible = |r: &Resource| -> bool {
2024                    !disabled.contains(&r.uri)
2025                        && self
2026                            .inner
2027                            .resource_filter
2028                            .as_ref()
2029                            .map(|f| f.is_visible(&self.session, r))
2030                            .unwrap_or(true)
2031                };
2032
2033                let mut resources: Vec<ResourceDefinition> = self
2034                    .inner
2035                    .resources
2036                    .values()
2037                    .filter(|r| is_visible(r))
2038                    .map(|r| r.definition())
2039                    .collect();
2040
2041                // Merge dynamic resources (static resources win on URI collision)
2042                #[cfg(feature = "dynamic-tools")]
2043                if let Some(ref dynamic) = self.inner.dynamic_resources {
2044                    let static_uris: HashSet<String> =
2045                        resources.iter().map(|r| r.uri.clone()).collect();
2046                    for r in dynamic.list() {
2047                        if !static_uris.contains(&r.uri) && is_visible(&r) {
2048                            resources.push(r.definition());
2049                        }
2050                    }
2051                }
2052
2053                resources.sort_by(|a, b| a.uri.cmp(&b.uri));
2054
2055                let (resources, next_cursor) =
2056                    paginate(resources, params.cursor.as_deref(), self.inner.page_size)?;
2057
2058                Ok(McpResponse::ListResources(ListResourcesResult {
2059                    resources,
2060                    next_cursor,
2061                    ttl_ms: self.inner.list_ttl_ms,
2062                    cache_scope: None,
2063                    meta: None,
2064                }))
2065            }
2066
2067            McpRequest::ListResourceTemplates(params) => {
2068                let mut resource_templates: Vec<ResourceTemplateDefinition> = self
2069                    .inner
2070                    .resource_templates
2071                    .iter()
2072                    .map(|t| t.definition())
2073                    .collect();
2074
2075                // Merge dynamic resource templates (static win on collision)
2076                #[cfg(feature = "dynamic-tools")]
2077                if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
2078                    let static_patterns: HashSet<String> = resource_templates
2079                        .iter()
2080                        .map(|t| t.uri_template.clone())
2081                        .collect();
2082                    for t in dynamic.list() {
2083                        if !static_patterns.contains(&t.uri_template) {
2084                            resource_templates.push(t.definition());
2085                        }
2086                    }
2087                }
2088
2089                resource_templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
2090
2091                let (resource_templates, next_cursor) = paginate(
2092                    resource_templates,
2093                    params.cursor.as_deref(),
2094                    self.inner.page_size,
2095                )?;
2096
2097                Ok(McpResponse::ListResourceTemplates(
2098                    ListResourceTemplatesResult {
2099                        resource_templates,
2100                        next_cursor,
2101                        ttl_ms: self.inner.list_ttl_ms,
2102                        cache_scope: None,
2103                        meta: None,
2104                    },
2105                ))
2106            }
2107
2108            McpRequest::ReadResource(params) => {
2109                // Disabled resources are reported as if they don't exist.
2110                if self
2111                    .inner
2112                    .disabled_resources
2113                    .read()
2114                    .unwrap()
2115                    .contains(&params.uri)
2116                {
2117                    return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2118                        &params.uri,
2119                    )));
2120                }
2121
2122                // First, try to find a static resource
2123                if let Some(resource) = self.inner.resources.get(&params.uri) {
2124                    // Check resource filter if configured
2125                    if let Some(filter) = &self.inner.resource_filter
2126                        && !filter.is_visible(&self.session, resource)
2127                    {
2128                        return Err(filter.denial_error(&params.uri));
2129                    }
2130
2131                    tracing::debug!(uri = %params.uri, "Reading static resource");
2132                    let ctx = self.create_context_with_extensions(request_id, None, &extensions);
2133                    let result = resource.read_with_context(ctx).await;
2134                    return Ok(McpResponse::ReadResource(result));
2135                }
2136
2137                // Try dynamic resources
2138                #[cfg(feature = "dynamic-tools")]
2139                #[allow(clippy::collapsible_if)]
2140                if let Some(ref dynamic) = self.inner.dynamic_resources {
2141                    if let Some(resource) = dynamic.get(&params.uri) {
2142                        if let Some(filter) = &self.inner.resource_filter
2143                            && !filter.is_visible(&self.session, &resource)
2144                        {
2145                            return Err(filter.denial_error(&params.uri));
2146                        }
2147                        tracing::debug!(uri = %params.uri, "Reading dynamic resource");
2148                        let ctx =
2149                            self.create_context_with_extensions(request_id, None, &extensions);
2150                        let result = resource.read_with_context(ctx).await;
2151                        return Ok(McpResponse::ReadResource(result));
2152                    }
2153                }
2154
2155                // Try static templates
2156                for template in &self.inner.resource_templates {
2157                    if let Some(variables) = template.match_uri(&params.uri) {
2158                        tracing::debug!(
2159                            uri = %params.uri,
2160                            template = %template.uri_template,
2161                            "Reading resource via template"
2162                        );
2163                        let result = template.read(&params.uri, variables).await?;
2164                        return Ok(McpResponse::ReadResource(result));
2165                    }
2166                }
2167
2168                // Try dynamic templates
2169                #[cfg(feature = "dynamic-tools")]
2170                #[allow(clippy::collapsible_if)]
2171                if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
2172                    if let Some((template, variables)) = dynamic.match_uri(&params.uri) {
2173                        tracing::debug!(
2174                            uri = %params.uri,
2175                            template = %template.uri_template,
2176                            "Reading resource via dynamic template"
2177                        );
2178                        let result = template.read(&params.uri, variables).await?;
2179                        return Ok(McpResponse::ReadResource(result));
2180                    }
2181                }
2182
2183                // No match found
2184                Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2185                    &params.uri,
2186                )))
2187            }
2188
2189            McpRequest::SubscribeResource(params) => {
2190                // Verify the resource exists
2191                if !self.inner.resources.contains_key(&params.uri) {
2192                    return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2193                        &params.uri,
2194                    )));
2195                }
2196
2197                tracing::debug!(uri = %params.uri, "Subscribing to resource");
2198                self.subscribe(&params.uri);
2199
2200                Ok(McpResponse::SubscribeResource(EmptyResult {}))
2201            }
2202
2203            McpRequest::UnsubscribeResource(params) => {
2204                // Verify the resource exists
2205                if !self.inner.resources.contains_key(&params.uri) {
2206                    return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
2207                        &params.uri,
2208                    )));
2209                }
2210
2211                tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
2212                self.unsubscribe(&params.uri);
2213
2214                Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
2215            }
2216
2217            McpRequest::ListPrompts(params) => {
2218                let disabled = self.inner.disabled_prompts.read().unwrap().clone();
2219                let is_visible = |p: &Prompt| -> bool {
2220                    !disabled.contains(&p.name)
2221                        && self
2222                            .inner
2223                            .prompt_filter
2224                            .as_ref()
2225                            .map(|f| f.is_visible(&self.session, p))
2226                            .unwrap_or(true)
2227                };
2228
2229                let mut prompts: Vec<PromptDefinition> = self
2230                    .inner
2231                    .prompts
2232                    .values()
2233                    .filter(|p| is_visible(p))
2234                    .map(|p| p.definition())
2235                    .collect();
2236
2237                // Merge dynamic prompts (static prompts win on name collision)
2238                #[cfg(feature = "dynamic-tools")]
2239                if let Some(ref dynamic) = self.inner.dynamic_prompts {
2240                    let static_names: HashSet<String> =
2241                        prompts.iter().map(|p| p.name.clone()).collect();
2242                    for p in dynamic.list() {
2243                        if !static_names.contains(&p.name) && is_visible(&p) {
2244                            prompts.push(p.definition());
2245                        }
2246                    }
2247                }
2248
2249                prompts.sort_by(|a, b| a.name.cmp(&b.name));
2250
2251                let (prompts, next_cursor) =
2252                    paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
2253
2254                Ok(McpResponse::ListPrompts(ListPromptsResult {
2255                    prompts,
2256                    next_cursor,
2257                    ttl_ms: self.inner.list_ttl_ms,
2258                    cache_scope: None,
2259                    meta: None,
2260                }))
2261            }
2262
2263            McpRequest::GetPrompt(params) => {
2264                // Disabled prompts are reported as if they don't exist.
2265                if self
2266                    .inner
2267                    .disabled_prompts
2268                    .read()
2269                    .unwrap()
2270                    .contains(&params.name)
2271                {
2272                    return Err(Error::JsonRpc(JsonRpcError::method_not_found(&format!(
2273                        "Prompt not found: {}",
2274                        params.name
2275                    ))));
2276                }
2277
2278                // Look up static prompts first, then dynamic
2279                let prompt = self.inner.prompts.get(&params.name).cloned();
2280                #[cfg(feature = "dynamic-tools")]
2281                let prompt = prompt.or_else(|| {
2282                    self.inner
2283                        .dynamic_prompts
2284                        .as_ref()
2285                        .and_then(|d| d.get(&params.name))
2286                });
2287                let prompt = prompt.ok_or_else(|| {
2288                    Error::JsonRpc(JsonRpcError::method_not_found(&format!(
2289                        "Prompt not found: {}",
2290                        params.name
2291                    )))
2292                })?;
2293
2294                // Check prompt filter if configured
2295                if let Some(filter) = &self.inner.prompt_filter
2296                    && !filter.is_visible(&self.session, &prompt)
2297                {
2298                    return Err(filter.denial_error(&params.name));
2299                }
2300
2301                tracing::debug!(name = %params.name, "Getting prompt");
2302                let ctx = self.create_context_with_extensions(request_id, None, &extensions);
2303                let result = prompt.get_with_context(ctx, params.arguments).await?;
2304
2305                Ok(McpResponse::GetPrompt(result))
2306            }
2307
2308            McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
2309
2310            McpRequest::ListTasks(params) => {
2311                let tasks = self.inner.task_store.list_tasks(params.status);
2312
2313                let (tasks, next_cursor) =
2314                    paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
2315
2316                Ok(McpResponse::ListTasks(ListTasksResult {
2317                    tasks,
2318                    next_cursor,
2319                }))
2320            }
2321
2322            McpRequest::GetTaskInfo(params) => {
2323                let task = self
2324                    .inner
2325                    .task_store
2326                    .get_task(&params.task_id)
2327                    .ok_or_else(|| {
2328                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2329                            "Task not found: {}",
2330                            params.task_id
2331                        )))
2332                    })?;
2333
2334                Ok(McpResponse::GetTaskInfo(task))
2335            }
2336
2337            McpRequest::GetTaskResult(params) => {
2338                // Wait for task to reach terminal state (blocks if still running)
2339                let (task_obj, result, error) = self
2340                    .inner
2341                    .task_store
2342                    .wait_for_completion(&params.task_id)
2343                    .await
2344                    .ok_or_else(|| {
2345                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2346                            "Task not found: {}",
2347                            params.task_id
2348                        )))
2349                    })?;
2350
2351                // Build _meta with related-task reference
2352                let meta = serde_json::json!({
2353                    "io.modelcontextprotocol/related-task": task_obj
2354                });
2355
2356                match task_obj.status {
2357                    TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
2358                        format!("Task {} was cancelled", params.task_id),
2359                    ))),
2360                    TaskStatus::Failed => {
2361                        let mut call_result = CallToolResult::error(
2362                            error.unwrap_or_else(|| "Task failed".to_string()),
2363                        );
2364                        call_result.meta = Some(meta);
2365                        Ok(McpResponse::GetTaskResult(call_result))
2366                    }
2367                    _ => {
2368                        let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
2369                        call_result.meta = Some(meta);
2370                        Ok(McpResponse::GetTaskResult(call_result))
2371                    }
2372                }
2373            }
2374
2375            McpRequest::UpdateTask(params) => {
2376                // SEP-2663 `tasks/update`: validate the task exists and
2377                // acknowledge with an empty result. tower-mcp does not yet
2378                // model server-initiated `inputRequests` for tasks (that's a
2379                // future MRTR-flavored feature), so we currently treat any
2380                // submitted `inputResponses` as ignorable per spec ("A server
2381                // SHOULD ignore any inputResponses mapped to a key that is
2382                // not currently outstanding").
2383                let _ = self
2384                    .inner
2385                    .task_store
2386                    .get_task(&params.task_id)
2387                    .ok_or_else(|| {
2388                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2389                            "Task not found: {}",
2390                            params.task_id
2391                        )))
2392                    })?;
2393                Ok(McpResponse::UpdateTask(EmptyResult {}))
2394            }
2395
2396            McpRequest::CancelTask(params) => {
2397                // First check if the task exists and is not already terminal
2398                let current = self
2399                    .inner
2400                    .task_store
2401                    .get_task(&params.task_id)
2402                    .ok_or_else(|| {
2403                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2404                            "Task not found: {}",
2405                            params.task_id
2406                        )))
2407                    })?;
2408
2409                if current.status.is_terminal() {
2410                    return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
2411                        "Task {} is already in terminal state: {}",
2412                        params.task_id, current.status
2413                    ))));
2414                }
2415
2416                let task_obj = self
2417                    .inner
2418                    .task_store
2419                    .cancel_task(&params.task_id, params.reason.as_deref())
2420                    .ok_or_else(|| {
2421                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2422                            "Task not found: {}",
2423                            params.task_id
2424                        )))
2425                    })?;
2426
2427                Ok(McpResponse::CancelTask(task_obj))
2428            }
2429
2430            McpRequest::SetLoggingLevel(params) => {
2431                tracing::debug!(level = ?params.level, "Client set logging level");
2432                if let Ok(mut level) = self.inner.min_log_level.write() {
2433                    *level = params.level;
2434                }
2435                Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
2436            }
2437
2438            McpRequest::Complete(params) => {
2439                tracing::debug!(
2440                    reference = ?params.reference,
2441                    argument = %params.argument.name,
2442                    "Completion request"
2443                );
2444
2445                // Delegate to registered completion handler if available
2446                if let Some(ref handler) = self.inner.completion_handler {
2447                    let result = handler(params).await?;
2448                    Ok(McpResponse::Complete(result))
2449                } else {
2450                    // No completion handler registered, return empty completions
2451                    Ok(McpResponse::Complete(CompleteResult::new(vec![])))
2452                }
2453            }
2454
2455            McpRequest::Unknown { method, .. } => {
2456                Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
2457            }
2458            _ => Err(Error::JsonRpc(JsonRpcError::method_not_found(
2459                "unknown method",
2460            ))),
2461        }
2462    }
2463
2464    /// Handle an MCP notification (no response expected)
2465    pub fn handle_notification(&self, notification: McpNotification) {
2466        match notification {
2467            McpNotification::Initialized => {
2468                let phase_before = self.session.phase();
2469                if self.session.mark_initialized() {
2470                    if phase_before == crate::session::SessionPhase::Uninitialized {
2471                        tracing::info!(
2472                            "Session initialized from uninitialized state (race resolved)"
2473                        );
2474                    } else {
2475                        tracing::info!("Session initialized, entering operation phase");
2476                    }
2477                } else {
2478                    tracing::warn!(
2479                        phase = ?self.session.phase(),
2480                        "Received initialized notification in unexpected state"
2481                    );
2482                }
2483            }
2484            McpNotification::Cancelled(params) => {
2485                if let Some(ref request_id) = params.request_id {
2486                    if self.cancel_request(request_id) {
2487                        tracing::info!(
2488                            request_id = ?request_id,
2489                            reason = ?params.reason,
2490                            "Request cancelled"
2491                        );
2492                    } else {
2493                        tracing::debug!(
2494                            request_id = ?request_id,
2495                            reason = ?params.reason,
2496                            "Cancellation requested for unknown request"
2497                        );
2498                    }
2499                } else {
2500                    tracing::debug!(
2501                        reason = ?params.reason,
2502                        "Cancellation notification received without request_id"
2503                    );
2504                }
2505            }
2506            McpNotification::Progress(params) => {
2507                tracing::trace!(
2508                    token = ?params.progress_token,
2509                    progress = params.progress,
2510                    total = ?params.total,
2511                    "Progress notification"
2512                );
2513                // Progress notifications from client are unusual but valid
2514            }
2515            McpNotification::RootsListChanged => {
2516                tracing::info!("Client roots list changed");
2517                // Server should re-request roots if needed
2518                // This is handled by the application layer
2519            }
2520            McpNotification::Unknown { method, .. } => {
2521                tracing::debug!(method = %method, "Unknown notification received");
2522            }
2523            _ => {
2524                tracing::debug!("Unrecognized notification variant received");
2525            }
2526        }
2527    }
2528}
2529
2530impl Default for McpRouter {
2531    fn default() -> Self {
2532        Self::new()
2533    }
2534}
2535
2536// =============================================================================
2537// Tower Service implementation
2538// =============================================================================
2539
2540// Re-export Extensions from context for backwards compatibility
2541pub use crate::context::Extensions;
2542
2543/// A map of tool names to their annotations, for use by middleware.
2544///
2545/// This is automatically inserted into [`RouterRequest::extensions`] for
2546/// `tools/call` requests, allowing middleware to inspect tool safety hints
2547/// (e.g., `read_only_hint`, `destructive_hint`) without needing direct
2548/// access to the router's tool registry.
2549///
2550/// # Example
2551///
2552/// ```rust,ignore
2553/// use tower_mcp::router::ToolAnnotationsMap;
2554/// use tower_mcp::protocol::McpRequest;
2555///
2556/// // In a middleware Service::call():
2557/// fn call(&mut self, req: RouterRequest) -> Self::Future {
2558///     if let McpRequest::CallTool(params) = &req.inner {
2559///         if let Some(map) = req.extensions.get::<ToolAnnotationsMap>() {
2560///             let annotations = map.get(&params.name);
2561///             // Check annotations.read_only_hint, destructive_hint, etc.
2562///         }
2563///     }
2564///     self.inner.call(req)
2565/// }
2566/// ```
2567#[derive(Debug, Clone)]
2568pub struct ToolAnnotationsMap {
2569    map: Arc<HashMap<String, ToolAnnotations>>,
2570}
2571
2572impl ToolAnnotationsMap {
2573    /// Look up annotations for a tool by name.
2574    ///
2575    /// Returns `None` if the tool has no annotations or doesn't exist.
2576    pub fn get(&self, tool_name: &str) -> Option<&ToolAnnotations> {
2577        self.map.get(tool_name)
2578    }
2579
2580    /// Check if a tool is read-only (does not modify state).
2581    ///
2582    /// Returns `false` if the tool has no annotations or doesn't exist
2583    /// (the MCP spec default for `readOnlyHint` is `false`).
2584    pub fn is_read_only(&self, tool_name: &str) -> bool {
2585        self.map.get(tool_name).is_some_and(|a| a.read_only_hint)
2586    }
2587
2588    /// Check if a tool may have destructive effects.
2589    ///
2590    /// Returns `true` if the tool has no annotations or doesn't exist
2591    /// (the MCP spec default for `destructiveHint` is `true`).
2592    pub fn is_destructive(&self, tool_name: &str) -> bool {
2593        self.map.get(tool_name).is_none_or(|a| a.destructive_hint)
2594    }
2595
2596    /// Check if a tool is idempotent.
2597    ///
2598    /// Returns `false` if the tool has no annotations or doesn't exist
2599    /// (the MCP spec default for `idempotentHint` is `false`).
2600    pub fn is_idempotent(&self, tool_name: &str) -> bool {
2601        self.map.get(tool_name).is_some_and(|a| a.idempotent_hint)
2602    }
2603}
2604
2605/// Request type for the tower Service implementation.
2606///
2607/// # Preserving extensions in middleware
2608///
2609/// When rewriting a request in middleware, use [`with_inner`](Self::with_inner)
2610/// or [`clone_with_inner`](Self::clone_with_inner) instead of constructing a
2611/// new `RouterRequest` directly. Constructing with `Extensions::new()` will
2612/// silently drop extensions set by earlier middleware layers (token claims,
2613/// RBAC context, etc.).
2614///
2615/// ```rust,ignore
2616/// // WRONG: drops extensions from earlier middleware
2617/// let rewritten = RouterRequest {
2618///     id: req.id.clone(),
2619///     inner: new_inner,
2620///     extensions: Extensions::new(),
2621/// };
2622///
2623/// // RIGHT: preserves extensions
2624/// let rewritten = req.with_inner(new_inner);
2625/// ```
2626#[derive(Debug, Clone)]
2627pub struct RouterRequest {
2628    /// The JSON-RPC request ID.
2629    pub id: RequestId,
2630    /// The parsed MCP request.
2631    pub inner: McpRequest,
2632    /// Type-map for passing data (e.g., `TokenClaims`) through middleware.
2633    pub extensions: Extensions,
2634}
2635
2636impl RouterRequest {
2637    /// Create a new `RouterRequest` with empty extensions.
2638    pub fn new(id: RequestId, inner: McpRequest) -> Self {
2639        Self {
2640            id,
2641            inner,
2642            extensions: Extensions::new(),
2643        }
2644    }
2645
2646    /// Replace the inner MCP request, preserving the id and extensions.
2647    ///
2648    /// This is the recommended way to rewrite requests in middleware,
2649    /// as it ensures extensions set by earlier middleware layers
2650    /// (e.g., token claims, RBAC context) are not lost.
2651    pub fn with_inner(self, inner: McpRequest) -> Self {
2652        Self {
2653            id: self.id,
2654            inner,
2655            extensions: self.extensions,
2656        }
2657    }
2658
2659    /// Replace both the id and inner MCP request, preserving extensions.
2660    ///
2661    /// Useful when middleware needs to assign a new request id
2662    /// (e.g., for fan-out or request duplication) while keeping
2663    /// the extensions from the original request.
2664    pub fn with_id_and_inner(self, id: RequestId, inner: McpRequest) -> Self {
2665        Self {
2666            id,
2667            inner,
2668            extensions: self.extensions,
2669        }
2670    }
2671
2672    /// Create a copy of this request with a different inner request,
2673    /// cloning the id and extensions from the original.
2674    ///
2675    /// Unlike [`with_inner`](Self::with_inner), this borrows `self`,
2676    /// which is useful when the original request is still needed
2677    /// (e.g., for traffic mirroring where you send the request to
2678    /// two backends).
2679    pub fn clone_with_inner(&self, inner: McpRequest) -> Self {
2680        Self {
2681            id: self.id.clone(),
2682            inner,
2683            extensions: self.extensions.clone(),
2684        }
2685    }
2686}
2687
2688/// Response type for the tower Service implementation
2689#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
2690pub struct RouterResponse {
2691    /// The JSON-RPC request ID this response corresponds to.
2692    pub id: RequestId,
2693    /// The MCP response or JSON-RPC error.
2694    pub inner: std::result::Result<McpResponse, JsonRpcError>,
2695}
2696
2697impl RouterResponse {
2698    /// Returns `true` if the response contains a JSON-RPC error.
2699    ///
2700    /// Since tower-mcp services use `Error = Infallible` (errors are carried
2701    /// inside the response, not in the `Result`), this method is useful for
2702    /// middleware that needs to inspect whether a request failed -- for example,
2703    /// retry or circuit breaker middleware.
2704    ///
2705    /// # Example
2706    ///
2707    /// ```rust,ignore
2708    /// // Response-based retry predicate for tower-resilience or similar
2709    /// fn is_retriable(response: &RouterResponse) -> bool {
2710    ///     response.is_error()
2711    /// }
2712    /// ```
2713    pub fn is_error(&self) -> bool {
2714        self.inner.is_err()
2715    }
2716
2717    /// Convert to JSON-RPC response
2718    pub fn into_jsonrpc(self) -> JsonRpcResponse {
2719        match self.inner {
2720            Ok(response) => match serde_json::to_value(response) {
2721                Ok(result) => JsonRpcResponse::result(self.id, result),
2722                Err(e) => {
2723                    tracing::error!(error = %e, "Failed to serialize response");
2724                    JsonRpcResponse::error(
2725                        Some(self.id),
2726                        JsonRpcError::internal_error(format!("Serialization error: {}", e)),
2727                    )
2728                }
2729            },
2730            Err(error) => JsonRpcResponse::error(Some(self.id), error),
2731        }
2732    }
2733}
2734
2735impl Service<RouterRequest> for McpRouter {
2736    type Response = RouterResponse;
2737    type Error = std::convert::Infallible; // Errors are in the response
2738    type Future =
2739        Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
2740
2741    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
2742        Poll::Ready(Ok(()))
2743    }
2744
2745    fn call(&mut self, req: RouterRequest) -> Self::Future {
2746        let router = self.clone();
2747        let request_id = req.id.clone();
2748        Box::pin(async move {
2749            let result = router.handle(req.id, req.inner, req.extensions).await;
2750            // Clean up tracking after request completes
2751            router.complete_request(&request_id);
2752            Ok(RouterResponse {
2753                id: request_id,
2754                // Map tower-mcp errors to JSON-RPC errors:
2755                // - Error::JsonRpc: forwarded as-is (preserves original code)
2756                // - Error::Tool: mapped to -32603 (Internal Error)
2757                // - All others: mapped to -32603 (Internal Error)
2758                inner: result.map_err(|e| match e {
2759                    Error::JsonRpc(err) => err,
2760                    Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
2761                    e => JsonRpcError::internal_error(e.to_string()),
2762                }),
2763            })
2764        })
2765    }
2766}
2767
2768#[cfg(test)]
2769mod tests {
2770    use super::*;
2771    use crate::extract::{Context, Json};
2772    use crate::jsonrpc::JsonRpcService;
2773    use crate::tool::ToolBuilder;
2774    use schemars::JsonSchema;
2775    use serde::Deserialize;
2776    use tower::ServiceExt;
2777
2778    #[derive(Debug, Deserialize, JsonSchema)]
2779    struct AddInput {
2780        a: i64,
2781        b: i64,
2782    }
2783
2784    /// Helper to initialize a router for testing
2785    async fn init_router(router: &mut McpRouter) {
2786        // Send initialize request
2787        let init_req = RouterRequest {
2788            id: RequestId::Number(0),
2789            inner: McpRequest::Initialize(InitializeParams {
2790                protocol_version: "2025-11-25".to_string(),
2791                capabilities: ClientCapabilities {
2792                    roots: None,
2793                    sampling: None,
2794                    elicitation: None,
2795                    tasks: None,
2796                    experimental: None,
2797                    extensions: None,
2798                },
2799                client_info: Implementation {
2800                    name: "test".to_string(),
2801                    version: "1.0".to_string(),
2802                    ..Default::default()
2803                },
2804                meta: None,
2805            }),
2806            extensions: Extensions::new(),
2807        };
2808        let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
2809        // Send initialized notification
2810        router.handle_notification(McpNotification::Initialized);
2811    }
2812
2813    #[tokio::test]
2814    async fn test_router_list_tools() {
2815        let add_tool = ToolBuilder::new("add")
2816            .description("Add two numbers")
2817            .handler(|input: AddInput| async move {
2818                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2819            })
2820            .build();
2821
2822        let mut router = McpRouter::new().tool(add_tool);
2823
2824        // Initialize session first
2825        init_router(&mut router).await;
2826
2827        let req = RouterRequest {
2828            id: RequestId::Number(1),
2829            inner: McpRequest::ListTools(ListToolsParams::default()),
2830            extensions: Extensions::new(),
2831        };
2832
2833        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2834
2835        match resp.inner {
2836            Ok(McpResponse::ListTools(result)) => {
2837                assert_eq!(result.tools.len(), 1);
2838                assert_eq!(result.tools[0].name, "add");
2839            }
2840            _ => panic!("Expected ListTools response"),
2841        }
2842    }
2843
2844    #[tokio::test]
2845    async fn test_router_call_tool() {
2846        let add_tool = ToolBuilder::new("add")
2847            .description("Add two numbers")
2848            .handler(|input: AddInput| async move {
2849                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2850            })
2851            .build();
2852
2853        let mut router = McpRouter::new().tool(add_tool);
2854
2855        // Initialize session first
2856        init_router(&mut router).await;
2857
2858        let req = RouterRequest {
2859            id: RequestId::Number(1),
2860            inner: McpRequest::CallTool(CallToolParams {
2861                name: "add".to_string(),
2862                arguments: serde_json::json!({"a": 2, "b": 3}),
2863                meta: None,
2864                task: None,
2865            }),
2866            extensions: Extensions::new(),
2867        };
2868
2869        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2870
2871        match resp.inner {
2872            Ok(McpResponse::CallTool(result)) => {
2873                assert!(!result.is_error);
2874                // Check the text content
2875                match &result.content[0] {
2876                    Content::Text { text, .. } => assert_eq!(text, "5"),
2877                    _ => panic!("Expected text content"),
2878                }
2879            }
2880            _ => panic!("Expected CallTool response"),
2881        }
2882    }
2883
2884    /// Helper to initialize a JsonRpcService for testing
2885    async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2886        let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2887            "protocolVersion": "2025-11-25",
2888            "capabilities": {},
2889            "clientInfo": { "name": "test", "version": "1.0" }
2890        }));
2891        let _ = service.call_single(init_req).await.unwrap();
2892        router.handle_notification(McpNotification::Initialized);
2893    }
2894
2895    #[tokio::test]
2896    async fn test_jsonrpc_service() {
2897        let add_tool = ToolBuilder::new("add")
2898            .description("Add two numbers")
2899            .handler(|input: AddInput| async move {
2900                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2901            })
2902            .build();
2903
2904        let router = McpRouter::new().tool(add_tool);
2905        let mut service = JsonRpcService::new(router.clone());
2906
2907        // Initialize session first
2908        init_jsonrpc_service(&mut service, &router).await;
2909
2910        let req = JsonRpcRequest::new(1, "tools/list");
2911
2912        let resp = service.call_single(req).await.unwrap();
2913
2914        match resp {
2915            JsonRpcResponse::Result(r) => {
2916                assert_eq!(r.id, RequestId::Number(1));
2917                let tools = r.result.get("tools").unwrap().as_array().unwrap();
2918                assert_eq!(tools.len(), 1);
2919            }
2920            JsonRpcResponse::Error(_) => panic!("Expected success response"),
2921            _ => panic!("unexpected response variant"),
2922        }
2923    }
2924
2925    #[tokio::test]
2926    async fn test_batch_request() {
2927        let add_tool = ToolBuilder::new("add")
2928            .description("Add two numbers")
2929            .handler(|input: AddInput| async move {
2930                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2931            })
2932            .build();
2933
2934        let router = McpRouter::new().tool(add_tool);
2935        let mut service = JsonRpcService::new(router.clone());
2936
2937        // Initialize session first
2938        init_jsonrpc_service(&mut service, &router).await;
2939
2940        // Create a batch of requests
2941        let requests = vec![
2942            JsonRpcRequest::new(1, "tools/list"),
2943            JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2944                "name": "add",
2945                "arguments": {"a": 10, "b": 20}
2946            })),
2947            JsonRpcRequest::new(3, "ping"),
2948        ];
2949
2950        let responses = service.call_batch(requests).await.unwrap();
2951
2952        assert_eq!(responses.len(), 3);
2953
2954        // Check first response (tools/list)
2955        match &responses[0] {
2956            JsonRpcResponse::Result(r) => {
2957                assert_eq!(r.id, RequestId::Number(1));
2958                let tools = r.result.get("tools").unwrap().as_array().unwrap();
2959                assert_eq!(tools.len(), 1);
2960            }
2961            JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2962            _ => panic!("unexpected response variant"),
2963        }
2964
2965        // Check second response (tools/call)
2966        match &responses[1] {
2967            JsonRpcResponse::Result(r) => {
2968                assert_eq!(r.id, RequestId::Number(2));
2969                let content = r.result.get("content").unwrap().as_array().unwrap();
2970                let text = content[0].get("text").unwrap().as_str().unwrap();
2971                assert_eq!(text, "30");
2972            }
2973            JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2974            _ => panic!("unexpected response variant"),
2975        }
2976
2977        // Check third response (ping)
2978        match &responses[2] {
2979            JsonRpcResponse::Result(r) => {
2980                assert_eq!(r.id, RequestId::Number(3));
2981            }
2982            JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2983            _ => panic!("unexpected response variant"),
2984        }
2985    }
2986
2987    #[tokio::test]
2988    async fn test_empty_batch_error() {
2989        let router = McpRouter::new();
2990        let mut service = JsonRpcService::new(router);
2991
2992        let result = service.call_batch(vec![]).await;
2993        assert!(result.is_err());
2994    }
2995
2996    // =========================================================================
2997    // Progress Token Tests
2998    // =========================================================================
2999
3000    #[tokio::test]
3001    async fn test_progress_token_extraction() {
3002        use crate::context::{ServerNotification, notification_channel};
3003        use crate::protocol::ProgressToken;
3004        use std::sync::Arc;
3005        use std::sync::atomic::{AtomicBool, Ordering};
3006
3007        // Track whether progress was reported
3008        let progress_reported = Arc::new(AtomicBool::new(false));
3009        let progress_ref = progress_reported.clone();
3010
3011        // Create a tool that reports progress
3012        let tool = ToolBuilder::new("progress_tool")
3013            .description("Tool that reports progress")
3014            .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
3015                let reported = progress_ref.clone();
3016                async move {
3017                    // Report progress - this should work if token was extracted
3018                    ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
3019                        .await;
3020                    reported.store(true, Ordering::SeqCst);
3021                    Ok(CallToolResult::text("done"))
3022                }
3023            })
3024            .build();
3025
3026        // Set up notification channel
3027        let (tx, mut rx) = notification_channel(10);
3028        let router = McpRouter::new().with_notification_sender(tx).tool(tool);
3029        let mut service = JsonRpcService::new(router.clone());
3030
3031        // Initialize
3032        init_jsonrpc_service(&mut service, &router).await;
3033
3034        // Call tool WITH progress token in _meta
3035        let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3036            "name": "progress_tool",
3037            "arguments": {"a": 1, "b": 2},
3038            "_meta": {
3039                "progressToken": "test-token-123"
3040            }
3041        }));
3042
3043        let resp = service.call_single(req).await.unwrap();
3044
3045        // Verify the tool was called successfully
3046        match resp {
3047            JsonRpcResponse::Result(_) => {}
3048            JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
3049            _ => panic!("unexpected response variant"),
3050        }
3051
3052        // Verify progress was reported by handler
3053        assert!(progress_reported.load(Ordering::SeqCst));
3054
3055        // Verify progress notification was sent through channel
3056        let notification = rx.try_recv().expect("Expected progress notification");
3057        match notification {
3058            ServerNotification::Progress(params) => {
3059                assert_eq!(
3060                    params.progress_token,
3061                    ProgressToken::String("test-token-123".to_string())
3062                );
3063                assert_eq!(params.progress, 50.0);
3064                assert_eq!(params.total, Some(100.0));
3065                assert_eq!(params.message.as_deref(), Some("Halfway"));
3066            }
3067            _ => panic!("Expected Progress notification"),
3068        }
3069    }
3070
3071    #[tokio::test]
3072    async fn test_tool_call_without_progress_token() {
3073        use crate::context::notification_channel;
3074        use std::sync::Arc;
3075        use std::sync::atomic::{AtomicBool, Ordering};
3076
3077        let progress_attempted = Arc::new(AtomicBool::new(false));
3078        let progress_ref = progress_attempted.clone();
3079
3080        let tool = ToolBuilder::new("no_token_tool")
3081            .description("Tool that tries to report progress without token")
3082            .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
3083                let attempted = progress_ref.clone();
3084                async move {
3085                    // Try to report progress - should be a no-op without token
3086                    ctx.report_progress(50.0, Some(100.0), None).await;
3087                    attempted.store(true, Ordering::SeqCst);
3088                    Ok(CallToolResult::text("done"))
3089                }
3090            })
3091            .build();
3092
3093        let (tx, mut rx) = notification_channel(10);
3094        let router = McpRouter::new().with_notification_sender(tx).tool(tool);
3095        let mut service = JsonRpcService::new(router.clone());
3096
3097        init_jsonrpc_service(&mut service, &router).await;
3098
3099        // Call tool WITHOUT progress token
3100        let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3101            "name": "no_token_tool",
3102            "arguments": {"a": 1, "b": 2}
3103        }));
3104
3105        let resp = service.call_single(req).await.unwrap();
3106        assert!(matches!(resp, JsonRpcResponse::Result(_)));
3107
3108        // Handler was called
3109        assert!(progress_attempted.load(Ordering::SeqCst));
3110
3111        // But no notification was sent (no progress token)
3112        assert!(rx.try_recv().is_err());
3113    }
3114
3115    #[tokio::test]
3116    async fn test_batch_errors_returned_not_dropped() {
3117        let add_tool = ToolBuilder::new("add")
3118            .description("Add two numbers")
3119            .handler(|input: AddInput| async move {
3120                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3121            })
3122            .build();
3123
3124        let router = McpRouter::new().tool(add_tool);
3125        let mut service = JsonRpcService::new(router.clone());
3126
3127        init_jsonrpc_service(&mut service, &router).await;
3128
3129        // Create a batch with one valid and one invalid request
3130        let requests = vec![
3131            // Valid request
3132            JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
3133                "name": "add",
3134                "arguments": {"a": 10, "b": 20}
3135            })),
3136            // Invalid request - tool doesn't exist
3137            JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
3138                "name": "nonexistent_tool",
3139                "arguments": {}
3140            })),
3141            // Another valid request
3142            JsonRpcRequest::new(3, "ping"),
3143        ];
3144
3145        let responses = service.call_batch(requests).await.unwrap();
3146
3147        // All three requests should have responses (errors are not dropped)
3148        assert_eq!(responses.len(), 3);
3149
3150        // First should be success
3151        match &responses[0] {
3152            JsonRpcResponse::Result(r) => {
3153                assert_eq!(r.id, RequestId::Number(1));
3154            }
3155            JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
3156            _ => panic!("unexpected response variant"),
3157        }
3158
3159        // Second should be an error (tool not found)
3160        match &responses[1] {
3161            JsonRpcResponse::Error(e) => {
3162                assert_eq!(e.id, Some(RequestId::Number(2)));
3163                // Error should indicate method not found
3164                assert!(e.error.message.contains("not found") || e.error.code == -32601);
3165            }
3166            JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
3167            _ => panic!("unexpected response variant"),
3168        }
3169
3170        // Third should be success
3171        match &responses[2] {
3172            JsonRpcResponse::Result(r) => {
3173                assert_eq!(r.id, RequestId::Number(3));
3174            }
3175            JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
3176            _ => panic!("unexpected response variant"),
3177        }
3178    }
3179
3180    // =========================================================================
3181    // Resource Template Tests
3182    // =========================================================================
3183
3184    #[tokio::test]
3185    async fn test_list_resource_templates() {
3186        use crate::resource::ResourceTemplateBuilder;
3187        use std::collections::HashMap;
3188
3189        let template = ResourceTemplateBuilder::new("file:///{path}")
3190            .name("Project Files")
3191            .description("Access project files")
3192            .handler(|uri: String, _vars: HashMap<String, String>| async move {
3193                Ok(ReadResourceResult {
3194                    contents: vec![ResourceContent {
3195                        uri,
3196                        mime_type: None,
3197                        text: None,
3198                        blob: None,
3199                        meta: None,
3200                    }],
3201                    meta: None,
3202                })
3203            });
3204
3205        let mut router = McpRouter::new().resource_template(template);
3206
3207        // Initialize session
3208        init_router(&mut router).await;
3209
3210        let req = RouterRequest {
3211            id: RequestId::Number(1),
3212            inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
3213            extensions: Extensions::new(),
3214        };
3215
3216        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3217
3218        match resp.inner {
3219            Ok(McpResponse::ListResourceTemplates(result)) => {
3220                assert_eq!(result.resource_templates.len(), 1);
3221                assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
3222                assert_eq!(result.resource_templates[0].name, "Project Files");
3223            }
3224            _ => panic!("Expected ListResourceTemplates response"),
3225        }
3226    }
3227
3228    #[tokio::test]
3229    async fn test_read_resource_via_template() {
3230        use crate::resource::ResourceTemplateBuilder;
3231        use std::collections::HashMap;
3232
3233        let template = ResourceTemplateBuilder::new("db://users/{id}")
3234            .name("User Records")
3235            .handler(|uri: String, vars: HashMap<String, String>| async move {
3236                let id = vars.get("id").unwrap().clone();
3237                Ok(ReadResourceResult {
3238                    contents: vec![ResourceContent {
3239                        uri,
3240                        mime_type: Some("application/json".to_string()),
3241                        text: Some(format!(r#"{{"id": "{}"}}"#, id)),
3242                        blob: None,
3243                        meta: None,
3244                    }],
3245                    meta: None,
3246                })
3247            });
3248
3249        let mut router = McpRouter::new().resource_template(template);
3250
3251        // Initialize session
3252        init_router(&mut router).await;
3253
3254        // Read a resource that matches the template
3255        let req = RouterRequest {
3256            id: RequestId::Number(1),
3257            inner: McpRequest::ReadResource(ReadResourceParams {
3258                uri: "db://users/123".to_string(),
3259                meta: None,
3260            }),
3261            extensions: Extensions::new(),
3262        };
3263
3264        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3265
3266        match resp.inner {
3267            Ok(McpResponse::ReadResource(result)) => {
3268                assert_eq!(result.contents.len(), 1);
3269                assert_eq!(result.contents[0].uri, "db://users/123");
3270                assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
3271            }
3272            _ => panic!("Expected ReadResource response"),
3273        }
3274    }
3275
3276    #[tokio::test]
3277    async fn test_static_resource_takes_precedence_over_template() {
3278        use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
3279        use std::collections::HashMap;
3280
3281        // Template that would match the same URI
3282        let template = ResourceTemplateBuilder::new("file:///{path}")
3283            .name("Files Template")
3284            .handler(|uri: String, _vars: HashMap<String, String>| async move {
3285                Ok(ReadResourceResult {
3286                    contents: vec![ResourceContent {
3287                        uri,
3288                        mime_type: None,
3289                        text: Some("from template".to_string()),
3290                        blob: None,
3291                        meta: None,
3292                    }],
3293                    meta: None,
3294                })
3295            });
3296
3297        // Static resource with exact URI
3298        let static_resource = ResourceBuilder::new("file:///README.md")
3299            .name("README")
3300            .text("from static resource");
3301
3302        let mut router = McpRouter::new()
3303            .resource_template(template)
3304            .resource(static_resource);
3305
3306        // Initialize session
3307        init_router(&mut router).await;
3308
3309        // Read the static resource - should NOT go through template
3310        let req = RouterRequest {
3311            id: RequestId::Number(1),
3312            inner: McpRequest::ReadResource(ReadResourceParams {
3313                uri: "file:///README.md".to_string(),
3314                meta: None,
3315            }),
3316            extensions: Extensions::new(),
3317        };
3318
3319        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3320
3321        match resp.inner {
3322            Ok(McpResponse::ReadResource(result)) => {
3323                // Should get static resource, not template
3324                assert_eq!(
3325                    result.contents[0].text.as_deref(),
3326                    Some("from static resource")
3327                );
3328            }
3329            _ => panic!("Expected ReadResource response"),
3330        }
3331    }
3332
3333    #[tokio::test]
3334    async fn test_resource_not_found_when_no_match() {
3335        use crate::resource::ResourceTemplateBuilder;
3336        use std::collections::HashMap;
3337
3338        let template = ResourceTemplateBuilder::new("db://users/{id}")
3339            .name("Users")
3340            .handler(|uri: String, _vars: HashMap<String, String>| async move {
3341                Ok(ReadResourceResult {
3342                    contents: vec![ResourceContent {
3343                        uri,
3344                        mime_type: None,
3345                        text: None,
3346                        blob: None,
3347                        meta: None,
3348                    }],
3349                    meta: None,
3350                })
3351            });
3352
3353        let mut router = McpRouter::new().resource_template(template);
3354
3355        // Initialize session
3356        init_router(&mut router).await;
3357
3358        // Try to read a URI that doesn't match any resource or template
3359        let req = RouterRequest {
3360            id: RequestId::Number(1),
3361            inner: McpRequest::ReadResource(ReadResourceParams {
3362                uri: "db://posts/123".to_string(),
3363                meta: None,
3364            }),
3365            extensions: Extensions::new(),
3366        };
3367
3368        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3369
3370        match resp.inner {
3371            Err(err) => {
3372                assert!(err.message.contains("not found"));
3373            }
3374            Ok(_) => panic!("Expected error for non-matching URI"),
3375        }
3376    }
3377
3378    #[tokio::test]
3379    async fn test_capabilities_include_resources_with_only_templates() {
3380        use crate::resource::ResourceTemplateBuilder;
3381        use std::collections::HashMap;
3382
3383        let template = ResourceTemplateBuilder::new("file:///{path}")
3384            .name("Files")
3385            .handler(|uri: String, _vars: HashMap<String, String>| async move {
3386                Ok(ReadResourceResult {
3387                    contents: vec![ResourceContent {
3388                        uri,
3389                        mime_type: None,
3390                        text: None,
3391                        blob: None,
3392                        meta: None,
3393                    }],
3394                    meta: None,
3395                })
3396            });
3397
3398        let mut router = McpRouter::new().resource_template(template);
3399
3400        // Send initialize request and check capabilities
3401        let init_req = RouterRequest {
3402            id: RequestId::Number(0),
3403            inner: McpRequest::Initialize(InitializeParams {
3404                protocol_version: "2025-11-25".to_string(),
3405                capabilities: ClientCapabilities {
3406                    roots: None,
3407                    sampling: None,
3408                    elicitation: None,
3409                    tasks: None,
3410                    experimental: None,
3411                    extensions: None,
3412                },
3413                client_info: Implementation {
3414                    name: "test".to_string(),
3415                    version: "1.0".to_string(),
3416                    ..Default::default()
3417                },
3418                meta: None,
3419            }),
3420            extensions: Extensions::new(),
3421        };
3422        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3423
3424        match resp.inner {
3425            Ok(McpResponse::Initialize(result)) => {
3426                // Should have resources capability even though only templates registered
3427                assert!(result.capabilities.resources.is_some());
3428            }
3429            _ => panic!("Expected Initialize response"),
3430        }
3431    }
3432
3433    // =========================================================================
3434    // Logging Notification Tests
3435    // =========================================================================
3436
3437    #[tokio::test]
3438    async fn test_log_sends_notification() {
3439        use crate::context::notification_channel;
3440
3441        let (tx, mut rx) = notification_channel(10);
3442        let router = McpRouter::new().with_notification_sender(tx);
3443
3444        // Send an info log
3445        let sent = router.log_info("Test message");
3446        assert!(sent);
3447
3448        // Should receive the notification
3449        let notification = rx.try_recv().unwrap();
3450        match notification {
3451            ServerNotification::LogMessage(params) => {
3452                assert_eq!(params.level, LogLevel::Info);
3453                let data = params.data;
3454                assert_eq!(
3455                    data.get("message").unwrap().as_str().unwrap(),
3456                    "Test message"
3457                );
3458            }
3459            _ => panic!("Expected LogMessage notification"),
3460        }
3461    }
3462
3463    #[tokio::test]
3464    async fn test_log_with_custom_params() {
3465        use crate::context::notification_channel;
3466
3467        let (tx, mut rx) = notification_channel(10);
3468        let router = McpRouter::new().with_notification_sender(tx);
3469
3470        // Send a custom log message
3471        let params = LoggingMessageParams::new(
3472            LogLevel::Error,
3473            serde_json::json!({
3474                "error": "Connection failed",
3475                "host": "localhost"
3476            }),
3477        )
3478        .with_logger("database");
3479
3480        let sent = router.log(params);
3481        assert!(sent);
3482
3483        let notification = rx.try_recv().unwrap();
3484        match notification {
3485            ServerNotification::LogMessage(params) => {
3486                assert_eq!(params.level, LogLevel::Error);
3487                assert_eq!(params.logger.as_deref(), Some("database"));
3488                let data = params.data;
3489                assert_eq!(
3490                    data.get("error").unwrap().as_str().unwrap(),
3491                    "Connection failed"
3492                );
3493            }
3494            _ => panic!("Expected LogMessage notification"),
3495        }
3496    }
3497
3498    #[tokio::test]
3499    async fn test_log_without_channel_returns_false() {
3500        // Router without notification channel
3501        let router = McpRouter::new();
3502
3503        // Should return false when no channel configured
3504        assert!(!router.log_info("Test"));
3505        assert!(!router.log_warning("Test"));
3506        assert!(!router.log_error("Test"));
3507        assert!(!router.log_debug("Test"));
3508    }
3509
3510    #[tokio::test]
3511    async fn test_logging_capability_with_channel() {
3512        use crate::context::notification_channel;
3513
3514        let (tx, _rx) = notification_channel(10);
3515        let mut router = McpRouter::new().with_notification_sender(tx);
3516
3517        // Initialize and check capabilities
3518        let init_req = RouterRequest {
3519            id: RequestId::Number(0),
3520            inner: McpRequest::Initialize(InitializeParams {
3521                protocol_version: "2025-11-25".to_string(),
3522                capabilities: ClientCapabilities {
3523                    roots: None,
3524                    sampling: None,
3525                    elicitation: None,
3526                    tasks: None,
3527                    experimental: None,
3528                    extensions: None,
3529                },
3530                client_info: Implementation {
3531                    name: "test".to_string(),
3532                    version: "1.0".to_string(),
3533                    ..Default::default()
3534                },
3535                meta: None,
3536            }),
3537            extensions: Extensions::new(),
3538        };
3539        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3540
3541        match resp.inner {
3542            Ok(McpResponse::Initialize(result)) => {
3543                // Should have logging capability when notification channel is set
3544                assert!(result.capabilities.logging.is_some());
3545            }
3546            _ => panic!("Expected Initialize response"),
3547        }
3548    }
3549
3550    #[tokio::test]
3551    async fn test_no_logging_capability_without_channel() {
3552        let mut router = McpRouter::new();
3553
3554        // Initialize and check capabilities
3555        let init_req = RouterRequest {
3556            id: RequestId::Number(0),
3557            inner: McpRequest::Initialize(InitializeParams {
3558                protocol_version: "2025-11-25".to_string(),
3559                capabilities: ClientCapabilities {
3560                    roots: None,
3561                    sampling: None,
3562                    elicitation: None,
3563                    tasks: None,
3564                    experimental: None,
3565                    extensions: None,
3566                },
3567                client_info: Implementation {
3568                    name: "test".to_string(),
3569                    version: "1.0".to_string(),
3570                    ..Default::default()
3571                },
3572                meta: None,
3573            }),
3574            extensions: Extensions::new(),
3575        };
3576        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3577
3578        match resp.inner {
3579            Ok(McpResponse::Initialize(result)) => {
3580                // Should NOT have logging capability without notification channel
3581                assert!(result.capabilities.logging.is_none());
3582            }
3583            _ => panic!("Expected Initialize response"),
3584        }
3585    }
3586
3587    // =========================================================================
3588    // Task Lifecycle Tests
3589    // =========================================================================
3590
3591    #[tokio::test]
3592    async fn test_create_task_via_call_tool() {
3593        let add_tool = ToolBuilder::new("add")
3594            .description("Add two numbers")
3595            .task_support(TaskSupportMode::Optional)
3596            .handler(|input: AddInput| async move {
3597                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3598            })
3599            .build();
3600
3601        let mut router = McpRouter::new().tool(add_tool);
3602        init_router(&mut router).await;
3603
3604        let req = RouterRequest {
3605            id: RequestId::Number(1),
3606            inner: McpRequest::CallTool(CallToolParams {
3607                name: "add".to_string(),
3608                arguments: serde_json::json!({"a": 5, "b": 10}),
3609                meta: None,
3610                task: Some(TaskRequestParams { ttl: None }),
3611            }),
3612            extensions: Extensions::new(),
3613        };
3614
3615        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3616
3617        match resp.inner {
3618            Ok(McpResponse::CreateTask(result)) => {
3619                assert!(result.task.task_id.starts_with("task-"));
3620                assert_eq!(result.task.status, TaskStatus::Working);
3621            }
3622            _ => panic!("Expected CreateTask response"),
3623        }
3624    }
3625
3626    #[tokio::test]
3627    async fn test_list_tasks_empty() {
3628        let mut router = McpRouter::new();
3629        init_router(&mut router).await;
3630
3631        let req = RouterRequest {
3632            id: RequestId::Number(1),
3633            inner: McpRequest::ListTasks(ListTasksParams::default()),
3634            extensions: Extensions::new(),
3635        };
3636
3637        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3638
3639        match resp.inner {
3640            Ok(McpResponse::ListTasks(result)) => {
3641                assert!(result.tasks.is_empty());
3642            }
3643            _ => panic!("Expected ListTasks response"),
3644        }
3645    }
3646
3647    #[tokio::test]
3648    async fn test_task_lifecycle_complete() {
3649        let add_tool = ToolBuilder::new("add")
3650            .description("Add two numbers")
3651            .task_support(TaskSupportMode::Optional)
3652            .handler(|input: AddInput| async move {
3653                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3654            })
3655            .build();
3656
3657        let mut router = McpRouter::new().tool(add_tool);
3658        init_router(&mut router).await;
3659
3660        // Create task via tools/call with task params
3661        let req = RouterRequest {
3662            id: RequestId::Number(1),
3663            inner: McpRequest::CallTool(CallToolParams {
3664                name: "add".to_string(),
3665                arguments: serde_json::json!({"a": 7, "b": 8}),
3666                meta: None,
3667                task: Some(TaskRequestParams { ttl: None }),
3668            }),
3669            extensions: Extensions::new(),
3670        };
3671
3672        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3673        let task_id = match resp.inner {
3674            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3675            _ => panic!("Expected CreateTask response"),
3676        };
3677
3678        // Wait for task to complete
3679        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3680
3681        // Get task result
3682        let req = RouterRequest {
3683            id: RequestId::Number(2),
3684            inner: McpRequest::GetTaskResult(GetTaskResultParams {
3685                task_id: task_id.clone(),
3686                meta: None,
3687            }),
3688            extensions: Extensions::new(),
3689        };
3690
3691        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3692
3693        match resp.inner {
3694            Ok(McpResponse::GetTaskResult(result)) => {
3695                // Result should have _meta with related-task
3696                assert!(result.meta.is_some());
3697                // Check the result content
3698                match &result.content[0] {
3699                    Content::Text { text, .. } => assert_eq!(text, "15"),
3700                    _ => panic!("Expected text content"),
3701                }
3702            }
3703            _ => panic!("Expected GetTaskResult response"),
3704        }
3705    }
3706
3707    #[tokio::test]
3708    async fn test_task_cancellation() {
3709        // Use a slow tool to test cancellation
3710        let slow_tool = ToolBuilder::new("slow")
3711            .description("Slow tool")
3712            .task_support(TaskSupportMode::Optional)
3713            .handler(|_input: serde_json::Value| async move {
3714                tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
3715                Ok(CallToolResult::text("done"))
3716            })
3717            .build();
3718
3719        let mut router = McpRouter::new().tool(slow_tool);
3720        init_router(&mut router).await;
3721
3722        // Create task
3723        let req = RouterRequest {
3724            id: RequestId::Number(1),
3725            inner: McpRequest::CallTool(CallToolParams {
3726                name: "slow".to_string(),
3727                arguments: serde_json::json!({}),
3728                meta: None,
3729                task: Some(TaskRequestParams { ttl: None }),
3730            }),
3731            extensions: Extensions::new(),
3732        };
3733
3734        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3735        let task_id = match resp.inner {
3736            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3737            _ => panic!("Expected CreateTask response"),
3738        };
3739
3740        // Cancel the task
3741        let req = RouterRequest {
3742            id: RequestId::Number(2),
3743            inner: McpRequest::CancelTask(CancelTaskParams {
3744                task_id: task_id.clone(),
3745                reason: Some("Test cancellation".to_string()),
3746                meta: None,
3747            }),
3748            extensions: Extensions::new(),
3749        };
3750
3751        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3752
3753        match resp.inner {
3754            Ok(McpResponse::CancelTask(task_obj)) => {
3755                assert_eq!(task_obj.status, TaskStatus::Cancelled);
3756            }
3757            _ => panic!("Expected CancelTask response"),
3758        }
3759    }
3760
3761    #[tokio::test]
3762    async fn test_get_task_info() {
3763        let add_tool = ToolBuilder::new("add")
3764            .description("Add two numbers")
3765            .task_support(TaskSupportMode::Optional)
3766            .handler(|input: AddInput| async move {
3767                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3768            })
3769            .build();
3770
3771        let mut router = McpRouter::new().tool(add_tool);
3772        init_router(&mut router).await;
3773
3774        // Create task with TTL
3775        let req = RouterRequest {
3776            id: RequestId::Number(1),
3777            inner: McpRequest::CallTool(CallToolParams {
3778                name: "add".to_string(),
3779                arguments: serde_json::json!({"a": 1, "b": 2}),
3780                meta: None,
3781                task: Some(TaskRequestParams { ttl: Some(600_000) }),
3782            }),
3783            extensions: Extensions::new(),
3784        };
3785
3786        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3787        let task_id = match resp.inner {
3788            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3789            _ => panic!("Expected CreateTask response"),
3790        };
3791
3792        // Get task info
3793        let req = RouterRequest {
3794            id: RequestId::Number(2),
3795            inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3796                task_id: task_id.clone(),
3797                meta: None,
3798            }),
3799            extensions: Extensions::new(),
3800        };
3801
3802        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3803
3804        match resp.inner {
3805            Ok(McpResponse::GetTaskInfo(info)) => {
3806                assert_eq!(info.task_id, task_id);
3807                assert!(info.created_at.contains('T')); // ISO 8601
3808                assert_eq!(info.ttl, Some(600_000));
3809            }
3810            _ => panic!("Expected GetTaskInfo response"),
3811        }
3812    }
3813
3814    #[tokio::test]
3815    async fn test_task_forbidden_tool_rejects_task_params() {
3816        let tool = ToolBuilder::new("sync_only")
3817            .description("Sync only tool")
3818            .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3819            .build();
3820
3821        let mut router = McpRouter::new().tool(tool);
3822        init_router(&mut router).await;
3823
3824        // Try to create task on a tool with Forbidden task support
3825        let req = RouterRequest {
3826            id: RequestId::Number(1),
3827            inner: McpRequest::CallTool(CallToolParams {
3828                name: "sync_only".to_string(),
3829                arguments: serde_json::json!({}),
3830                meta: None,
3831                task: Some(TaskRequestParams { ttl: None }),
3832            }),
3833            extensions: Extensions::new(),
3834        };
3835
3836        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3837
3838        match resp.inner {
3839            Err(e) => {
3840                assert!(e.message.contains("does not support async tasks"));
3841            }
3842            _ => panic!("Expected error response"),
3843        }
3844    }
3845
3846    #[tokio::test]
3847    async fn test_get_nonexistent_task() {
3848        let mut router = McpRouter::new();
3849        init_router(&mut router).await;
3850
3851        let req = RouterRequest {
3852            id: RequestId::Number(1),
3853            inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3854                task_id: "task-999".to_string(),
3855                meta: None,
3856            }),
3857            extensions: Extensions::new(),
3858        };
3859
3860        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3861
3862        match resp.inner {
3863            Err(e) => {
3864                assert!(e.message.contains("not found"));
3865            }
3866            _ => panic!("Expected error response"),
3867        }
3868    }
3869
3870    // =========================================================================
3871    // Resource Subscription Tests
3872    // =========================================================================
3873
3874    #[tokio::test]
3875    async fn test_subscribe_to_resource() {
3876        use crate::resource::ResourceBuilder;
3877
3878        let resource = ResourceBuilder::new("file:///test.txt")
3879            .name("Test File")
3880            .text("Hello");
3881
3882        let mut router = McpRouter::new().resource(resource);
3883        init_router(&mut router).await;
3884
3885        // Subscribe to the resource
3886        let req = RouterRequest {
3887            id: RequestId::Number(1),
3888            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3889                uri: "file:///test.txt".to_string(),
3890                meta: None,
3891            }),
3892            extensions: Extensions::new(),
3893        };
3894
3895        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3896
3897        match resp.inner {
3898            Ok(McpResponse::SubscribeResource(_)) => {
3899                // Should be subscribed now
3900                assert!(router.is_subscribed("file:///test.txt"));
3901            }
3902            _ => panic!("Expected SubscribeResource response"),
3903        }
3904    }
3905
3906    #[tokio::test]
3907    async fn test_unsubscribe_from_resource() {
3908        use crate::resource::ResourceBuilder;
3909
3910        let resource = ResourceBuilder::new("file:///test.txt")
3911            .name("Test File")
3912            .text("Hello");
3913
3914        let mut router = McpRouter::new().resource(resource);
3915        init_router(&mut router).await;
3916
3917        // Subscribe first
3918        let req = RouterRequest {
3919            id: RequestId::Number(1),
3920            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3921                uri: "file:///test.txt".to_string(),
3922                meta: None,
3923            }),
3924            extensions: Extensions::new(),
3925        };
3926        let _ = router.ready().await.unwrap().call(req).await.unwrap();
3927        assert!(router.is_subscribed("file:///test.txt"));
3928
3929        // Now unsubscribe
3930        let req = RouterRequest {
3931            id: RequestId::Number(2),
3932            inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3933                uri: "file:///test.txt".to_string(),
3934                meta: None,
3935            }),
3936            extensions: Extensions::new(),
3937        };
3938
3939        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3940
3941        match resp.inner {
3942            Ok(McpResponse::UnsubscribeResource(_)) => {
3943                // Should no longer be subscribed
3944                assert!(!router.is_subscribed("file:///test.txt"));
3945            }
3946            _ => panic!("Expected UnsubscribeResource response"),
3947        }
3948    }
3949
3950    #[tokio::test]
3951    async fn test_subscribe_nonexistent_resource() {
3952        let mut router = McpRouter::new();
3953        init_router(&mut router).await;
3954
3955        let req = RouterRequest {
3956            id: RequestId::Number(1),
3957            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3958                uri: "file:///nonexistent.txt".to_string(),
3959                meta: None,
3960            }),
3961            extensions: Extensions::new(),
3962        };
3963
3964        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3965
3966        match resp.inner {
3967            Err(e) => {
3968                assert!(e.message.contains("not found"));
3969            }
3970            _ => panic!("Expected error response"),
3971        }
3972    }
3973
3974    #[tokio::test]
3975    async fn test_notify_resource_updated() {
3976        use crate::context::notification_channel;
3977        use crate::resource::ResourceBuilder;
3978
3979        let (tx, mut rx) = notification_channel(10);
3980
3981        let resource = ResourceBuilder::new("file:///test.txt")
3982            .name("Test File")
3983            .text("Hello");
3984
3985        let router = McpRouter::new()
3986            .resource(resource)
3987            .with_notification_sender(tx);
3988
3989        // First, manually subscribe (simulate subscription)
3990        router.subscribe("file:///test.txt");
3991
3992        // Now notify
3993        let sent = router.notify_resource_updated("file:///test.txt");
3994        assert!(sent);
3995
3996        // Check the notification was sent
3997        let notification = rx.try_recv().unwrap();
3998        match notification {
3999            ServerNotification::ResourceUpdated { uri } => {
4000                assert_eq!(uri, "file:///test.txt");
4001            }
4002            _ => panic!("Expected ResourceUpdated notification"),
4003        }
4004    }
4005
4006    #[tokio::test]
4007    async fn test_notify_resource_updated_not_subscribed() {
4008        use crate::context::notification_channel;
4009        use crate::resource::ResourceBuilder;
4010
4011        let (tx, mut rx) = notification_channel(10);
4012
4013        let resource = ResourceBuilder::new("file:///test.txt")
4014            .name("Test File")
4015            .text("Hello");
4016
4017        let router = McpRouter::new()
4018            .resource(resource)
4019            .with_notification_sender(tx);
4020
4021        // Try to notify without subscribing
4022        let sent = router.notify_resource_updated("file:///test.txt");
4023        assert!(!sent); // Should not send because not subscribed
4024
4025        // Channel should be empty
4026        assert!(rx.try_recv().is_err());
4027    }
4028
4029    #[tokio::test]
4030    async fn test_notify_resources_list_changed() {
4031        use crate::context::notification_channel;
4032
4033        let (tx, mut rx) = notification_channel(10);
4034        let router = McpRouter::new().with_notification_sender(tx);
4035
4036        let sent = router.notify_resources_list_changed();
4037        assert!(sent);
4038
4039        let notification = rx.try_recv().unwrap();
4040        match notification {
4041            ServerNotification::ResourcesListChanged => {}
4042            _ => panic!("Expected ResourcesListChanged notification"),
4043        }
4044    }
4045
4046    #[tokio::test]
4047    async fn test_subscribed_uris() {
4048        use crate::resource::ResourceBuilder;
4049
4050        let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
4051
4052        let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
4053
4054        let router = McpRouter::new().resource(resource1).resource(resource2);
4055
4056        // Subscribe to both
4057        router.subscribe("file:///a.txt");
4058        router.subscribe("file:///b.txt");
4059
4060        let uris = router.subscribed_uris();
4061        assert_eq!(uris.len(), 2);
4062        assert!(uris.contains(&"file:///a.txt".to_string()));
4063        assert!(uris.contains(&"file:///b.txt".to_string()));
4064    }
4065
4066    #[tokio::test]
4067    async fn test_subscription_capability_advertised() {
4068        use crate::resource::ResourceBuilder;
4069
4070        let resource = ResourceBuilder::new("file:///test.txt")
4071            .name("Test")
4072            .text("Hello");
4073
4074        let mut router = McpRouter::new().resource(resource);
4075
4076        // Initialize and check capabilities
4077        let init_req = RouterRequest {
4078            id: RequestId::Number(0),
4079            inner: McpRequest::Initialize(InitializeParams {
4080                protocol_version: "2025-11-25".to_string(),
4081                capabilities: ClientCapabilities {
4082                    roots: None,
4083                    sampling: None,
4084                    elicitation: None,
4085                    tasks: None,
4086                    experimental: None,
4087                    extensions: None,
4088                },
4089                client_info: Implementation {
4090                    name: "test".to_string(),
4091                    version: "1.0".to_string(),
4092                    ..Default::default()
4093                },
4094                meta: None,
4095            }),
4096            extensions: Extensions::new(),
4097        };
4098        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
4099
4100        match resp.inner {
4101            Ok(McpResponse::Initialize(result)) => {
4102                // Should have resources capability with subscribe enabled
4103                let resources_cap = result.capabilities.resources.unwrap();
4104                assert!(resources_cap.subscribe);
4105            }
4106            _ => panic!("Expected Initialize response"),
4107        }
4108    }
4109
4110    #[tokio::test]
4111    async fn test_completion_handler() {
4112        let router = McpRouter::new()
4113            .server_info("test", "1.0")
4114            .completion_handler(|params: CompleteParams| async move {
4115                // Return suggestions based on the argument value
4116                let prefix = &params.argument.value;
4117                let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
4118                    .into_iter()
4119                    .filter(|s| s.starts_with(prefix))
4120                    .map(String::from)
4121                    .collect();
4122                Ok(CompleteResult::new(suggestions))
4123            });
4124
4125        // Initialize
4126        let init_req = RouterRequest {
4127            id: RequestId::Number(0),
4128            inner: McpRequest::Initialize(InitializeParams {
4129                protocol_version: "2025-11-25".to_string(),
4130                capabilities: ClientCapabilities::default(),
4131                client_info: Implementation {
4132                    name: "test".to_string(),
4133                    version: "1.0".to_string(),
4134                    ..Default::default()
4135                },
4136                meta: None,
4137            }),
4138            extensions: Extensions::new(),
4139        };
4140        let resp = router
4141            .clone()
4142            .ready()
4143            .await
4144            .unwrap()
4145            .call(init_req)
4146            .await
4147            .unwrap();
4148
4149        // Check that completions capability is advertised
4150        match resp.inner {
4151            Ok(McpResponse::Initialize(result)) => {
4152                assert!(result.capabilities.completions.is_some());
4153            }
4154            _ => panic!("Expected Initialize response"),
4155        }
4156
4157        // Send initialized notification
4158        router.handle_notification(McpNotification::Initialized);
4159
4160        // Test completion request
4161        let complete_req = RouterRequest {
4162            id: RequestId::Number(1),
4163            inner: McpRequest::Complete(CompleteParams {
4164                reference: CompletionReference::prompt("test-prompt"),
4165                argument: CompletionArgument::new("query", "al"),
4166                context: None,
4167                meta: None,
4168            }),
4169            extensions: Extensions::new(),
4170        };
4171        let resp = router
4172            .clone()
4173            .ready()
4174            .await
4175            .unwrap()
4176            .call(complete_req)
4177            .await
4178            .unwrap();
4179
4180        match resp.inner {
4181            Ok(McpResponse::Complete(result)) => {
4182                assert_eq!(result.completion.values, vec!["alpha"]);
4183            }
4184            _ => panic!("Expected Complete response"),
4185        }
4186    }
4187
4188    #[tokio::test]
4189    async fn test_completion_without_handler_returns_empty() {
4190        let router = McpRouter::new().server_info("test", "1.0");
4191
4192        // Initialize
4193        let init_req = RouterRequest {
4194            id: RequestId::Number(0),
4195            inner: McpRequest::Initialize(InitializeParams {
4196                protocol_version: "2025-11-25".to_string(),
4197                capabilities: ClientCapabilities::default(),
4198                client_info: Implementation {
4199                    name: "test".to_string(),
4200                    version: "1.0".to_string(),
4201                    ..Default::default()
4202                },
4203                meta: None,
4204            }),
4205            extensions: Extensions::new(),
4206        };
4207        let resp = router
4208            .clone()
4209            .ready()
4210            .await
4211            .unwrap()
4212            .call(init_req)
4213            .await
4214            .unwrap();
4215
4216        // Check that completions capability is NOT advertised
4217        match resp.inner {
4218            Ok(McpResponse::Initialize(result)) => {
4219                assert!(result.capabilities.completions.is_none());
4220            }
4221            _ => panic!("Expected Initialize response"),
4222        }
4223
4224        // Send initialized notification
4225        router.handle_notification(McpNotification::Initialized);
4226
4227        // Test completion request still works but returns empty
4228        let complete_req = RouterRequest {
4229            id: RequestId::Number(1),
4230            inner: McpRequest::Complete(CompleteParams {
4231                reference: CompletionReference::prompt("test-prompt"),
4232                argument: CompletionArgument::new("query", "al"),
4233                context: None,
4234                meta: None,
4235            }),
4236            extensions: Extensions::new(),
4237        };
4238        let resp = router
4239            .clone()
4240            .ready()
4241            .await
4242            .unwrap()
4243            .call(complete_req)
4244            .await
4245            .unwrap();
4246
4247        match resp.inner {
4248            Ok(McpResponse::Complete(result)) => {
4249                assert!(result.completion.values.is_empty());
4250            }
4251            _ => panic!("Expected Complete response"),
4252        }
4253    }
4254
4255    #[tokio::test]
4256    async fn test_tool_filter_list() {
4257        use crate::filter::CapabilityFilter;
4258        use crate::tool::Tool;
4259
4260        let public_tool = ToolBuilder::new("public")
4261            .description("Public tool")
4262            .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
4263            .build();
4264
4265        let admin_tool = ToolBuilder::new("admin")
4266            .description("Admin tool")
4267            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4268            .build();
4269
4270        let mut router = McpRouter::new()
4271            .tool(public_tool)
4272            .tool(admin_tool)
4273            .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
4274
4275        // Initialize session
4276        init_router(&mut router).await;
4277
4278        let req = RouterRequest {
4279            id: RequestId::Number(1),
4280            inner: McpRequest::ListTools(ListToolsParams::default()),
4281            extensions: Extensions::new(),
4282        };
4283
4284        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4285
4286        match resp.inner {
4287            Ok(McpResponse::ListTools(result)) => {
4288                // Only public tool should be visible
4289                assert_eq!(result.tools.len(), 1);
4290                assert_eq!(result.tools[0].name, "public");
4291            }
4292            _ => panic!("Expected ListTools response"),
4293        }
4294    }
4295
4296    #[tokio::test]
4297    async fn test_tool_filter_call_denied() {
4298        use crate::filter::CapabilityFilter;
4299        use crate::tool::Tool;
4300
4301        let admin_tool = ToolBuilder::new("admin")
4302            .description("Admin tool")
4303            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4304            .build();
4305
4306        let mut router = McpRouter::new()
4307            .tool(admin_tool)
4308            .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); // Deny all
4309
4310        // Initialize session
4311        init_router(&mut router).await;
4312
4313        let req = RouterRequest {
4314            id: RequestId::Number(1),
4315            inner: McpRequest::CallTool(CallToolParams {
4316                name: "admin".to_string(),
4317                arguments: serde_json::json!({"a": 1, "b": 2}),
4318                meta: None,
4319                task: None,
4320            }),
4321            extensions: Extensions::new(),
4322        };
4323
4324        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4325
4326        // Should get method not found error (default denial behavior)
4327        match resp.inner {
4328            Err(e) => {
4329                assert_eq!(e.code, -32601); // Method not found
4330            }
4331            _ => panic!("Expected JsonRpc error"),
4332        }
4333    }
4334
4335    #[tokio::test]
4336    async fn test_tool_filter_call_allowed() {
4337        use crate::filter::CapabilityFilter;
4338        use crate::tool::Tool;
4339
4340        let public_tool = ToolBuilder::new("public")
4341            .description("Public tool")
4342            .handler(|input: AddInput| async move {
4343                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
4344            })
4345            .build();
4346
4347        let mut router = McpRouter::new()
4348            .tool(public_tool)
4349            .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); // Allow all
4350
4351        // Initialize session
4352        init_router(&mut router).await;
4353
4354        let req = RouterRequest {
4355            id: RequestId::Number(1),
4356            inner: McpRequest::CallTool(CallToolParams {
4357                name: "public".to_string(),
4358                arguments: serde_json::json!({"a": 1, "b": 2}),
4359                meta: None,
4360                task: None,
4361            }),
4362            extensions: Extensions::new(),
4363        };
4364
4365        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4366
4367        match resp.inner {
4368            Ok(McpResponse::CallTool(result)) => {
4369                assert!(!result.is_error);
4370            }
4371            _ => panic!("Expected CallTool response"),
4372        }
4373    }
4374
4375    #[tokio::test]
4376    async fn test_tool_filter_custom_denial() {
4377        use crate::filter::{CapabilityFilter, DenialBehavior};
4378        use crate::tool::Tool;
4379
4380        let admin_tool = ToolBuilder::new("admin")
4381            .description("Admin tool")
4382            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
4383            .build();
4384
4385        let mut router = McpRouter::new().tool(admin_tool).tool_filter(
4386            CapabilityFilter::new(|_, _: &Tool| false)
4387                .denial_behavior(DenialBehavior::Unauthorized),
4388        );
4389
4390        // Initialize session
4391        init_router(&mut router).await;
4392
4393        let req = RouterRequest {
4394            id: RequestId::Number(1),
4395            inner: McpRequest::CallTool(CallToolParams {
4396                name: "admin".to_string(),
4397                arguments: serde_json::json!({"a": 1, "b": 2}),
4398                meta: None,
4399                task: None,
4400            }),
4401            extensions: Extensions::new(),
4402        };
4403
4404        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4405
4406        // Should get forbidden error
4407        match resp.inner {
4408            Err(e) => {
4409                assert_eq!(e.code, -32007); // Forbidden
4410                assert!(e.message.contains("Unauthorized"));
4411            }
4412            _ => panic!("Expected JsonRpc error"),
4413        }
4414    }
4415
4416    #[tokio::test]
4417    async fn test_resource_filter_list() {
4418        use crate::filter::CapabilityFilter;
4419        use crate::resource::{Resource, ResourceBuilder};
4420
4421        let public_resource = ResourceBuilder::new("file:///public.txt")
4422            .name("Public File")
4423            .text("public content");
4424
4425        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4426            .name("Secret File")
4427            .text("secret content");
4428
4429        let mut router = McpRouter::new()
4430            .resource(public_resource)
4431            .resource(secret_resource)
4432            .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
4433                !r.name.contains("Secret")
4434            }));
4435
4436        // Initialize session
4437        init_router(&mut router).await;
4438
4439        let req = RouterRequest {
4440            id: RequestId::Number(1),
4441            inner: McpRequest::ListResources(ListResourcesParams::default()),
4442            extensions: Extensions::new(),
4443        };
4444
4445        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4446
4447        match resp.inner {
4448            Ok(McpResponse::ListResources(result)) => {
4449                // Should only see public resource
4450                assert_eq!(result.resources.len(), 1);
4451                assert_eq!(result.resources[0].name, "Public File");
4452            }
4453            _ => panic!("Expected ListResources response"),
4454        }
4455    }
4456
4457    #[tokio::test]
4458    async fn test_resource_filter_read_denied() {
4459        use crate::filter::CapabilityFilter;
4460        use crate::resource::{Resource, ResourceBuilder};
4461
4462        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4463            .name("Secret File")
4464            .text("secret content");
4465
4466        let mut router = McpRouter::new()
4467            .resource(secret_resource)
4468            .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); // Deny all
4469
4470        // Initialize session
4471        init_router(&mut router).await;
4472
4473        let req = RouterRequest {
4474            id: RequestId::Number(1),
4475            inner: McpRequest::ReadResource(ReadResourceParams {
4476                uri: "file:///secret.txt".to_string(),
4477                meta: None,
4478            }),
4479            extensions: Extensions::new(),
4480        };
4481
4482        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4483
4484        // Should get method not found error (default denial behavior)
4485        match resp.inner {
4486            Err(e) => {
4487                assert_eq!(e.code, -32601); // Method not found
4488            }
4489            _ => panic!("Expected JsonRpc error"),
4490        }
4491    }
4492
4493    #[tokio::test]
4494    async fn test_resource_filter_read_allowed() {
4495        use crate::filter::CapabilityFilter;
4496        use crate::resource::{Resource, ResourceBuilder};
4497
4498        let public_resource = ResourceBuilder::new("file:///public.txt")
4499            .name("Public File")
4500            .text("public content");
4501
4502        let mut router = McpRouter::new()
4503            .resource(public_resource)
4504            .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); // Allow all
4505
4506        // Initialize session
4507        init_router(&mut router).await;
4508
4509        let req = RouterRequest {
4510            id: RequestId::Number(1),
4511            inner: McpRequest::ReadResource(ReadResourceParams {
4512                uri: "file:///public.txt".to_string(),
4513                meta: None,
4514            }),
4515            extensions: Extensions::new(),
4516        };
4517
4518        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4519
4520        match resp.inner {
4521            Ok(McpResponse::ReadResource(result)) => {
4522                assert_eq!(result.contents.len(), 1);
4523                assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
4524            }
4525            _ => panic!("Expected ReadResource response"),
4526        }
4527    }
4528
4529    #[tokio::test]
4530    async fn test_resource_filter_custom_denial() {
4531        use crate::filter::{CapabilityFilter, DenialBehavior};
4532        use crate::resource::{Resource, ResourceBuilder};
4533
4534        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4535            .name("Secret File")
4536            .text("secret content");
4537
4538        let mut router = McpRouter::new().resource(secret_resource).resource_filter(
4539            CapabilityFilter::new(|_, _: &Resource| false)
4540                .denial_behavior(DenialBehavior::Unauthorized),
4541        );
4542
4543        // Initialize session
4544        init_router(&mut router).await;
4545
4546        let req = RouterRequest {
4547            id: RequestId::Number(1),
4548            inner: McpRequest::ReadResource(ReadResourceParams {
4549                uri: "file:///secret.txt".to_string(),
4550                meta: None,
4551            }),
4552            extensions: Extensions::new(),
4553        };
4554
4555        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4556
4557        // Should get forbidden error
4558        match resp.inner {
4559            Err(e) => {
4560                assert_eq!(e.code, -32007); // Forbidden
4561                assert!(e.message.contains("Unauthorized"));
4562            }
4563            _ => panic!("Expected JsonRpc error"),
4564        }
4565    }
4566
4567    #[tokio::test]
4568    async fn test_prompt_filter_list() {
4569        use crate::filter::CapabilityFilter;
4570        use crate::prompt::{Prompt, PromptBuilder};
4571
4572        let public_prompt = PromptBuilder::new("greeting")
4573            .description("A greeting")
4574            .user_message("Hello!");
4575
4576        let admin_prompt = PromptBuilder::new("system_debug")
4577            .description("Admin prompt")
4578            .user_message("Debug");
4579
4580        let mut router = McpRouter::new()
4581            .prompt(public_prompt)
4582            .prompt(admin_prompt)
4583            .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
4584                !p.name.contains("system")
4585            }));
4586
4587        // Initialize session
4588        init_router(&mut router).await;
4589
4590        let req = RouterRequest {
4591            id: RequestId::Number(1),
4592            inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4593            extensions: Extensions::new(),
4594        };
4595
4596        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4597
4598        match resp.inner {
4599            Ok(McpResponse::ListPrompts(result)) => {
4600                // Should only see public prompt
4601                assert_eq!(result.prompts.len(), 1);
4602                assert_eq!(result.prompts[0].name, "greeting");
4603            }
4604            _ => panic!("Expected ListPrompts response"),
4605        }
4606    }
4607
4608    #[tokio::test]
4609    async fn test_prompt_filter_get_denied() {
4610        use crate::filter::CapabilityFilter;
4611        use crate::prompt::{Prompt, PromptBuilder};
4612        use std::collections::HashMap;
4613
4614        let admin_prompt = PromptBuilder::new("system_debug")
4615            .description("Admin prompt")
4616            .user_message("Debug");
4617
4618        let mut router = McpRouter::new()
4619            .prompt(admin_prompt)
4620            .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); // Deny all
4621
4622        // Initialize session
4623        init_router(&mut router).await;
4624
4625        let req = RouterRequest {
4626            id: RequestId::Number(1),
4627            inner: McpRequest::GetPrompt(GetPromptParams {
4628                name: "system_debug".to_string(),
4629                arguments: HashMap::new(),
4630                meta: None,
4631            }),
4632            extensions: Extensions::new(),
4633        };
4634
4635        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4636
4637        // Should get method not found error (default denial behavior)
4638        match resp.inner {
4639            Err(e) => {
4640                assert_eq!(e.code, -32601); // Method not found
4641            }
4642            _ => panic!("Expected JsonRpc error"),
4643        }
4644    }
4645
4646    #[tokio::test]
4647    async fn test_prompt_filter_get_allowed() {
4648        use crate::filter::CapabilityFilter;
4649        use crate::prompt::{Prompt, PromptBuilder};
4650        use std::collections::HashMap;
4651
4652        let public_prompt = PromptBuilder::new("greeting")
4653            .description("A greeting")
4654            .user_message("Hello!");
4655
4656        let mut router = McpRouter::new()
4657            .prompt(public_prompt)
4658            .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); // Allow all
4659
4660        // Initialize session
4661        init_router(&mut router).await;
4662
4663        let req = RouterRequest {
4664            id: RequestId::Number(1),
4665            inner: McpRequest::GetPrompt(GetPromptParams {
4666                name: "greeting".to_string(),
4667                arguments: HashMap::new(),
4668                meta: None,
4669            }),
4670            extensions: Extensions::new(),
4671        };
4672
4673        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4674
4675        match resp.inner {
4676            Ok(McpResponse::GetPrompt(result)) => {
4677                assert_eq!(result.messages.len(), 1);
4678            }
4679            _ => panic!("Expected GetPrompt response"),
4680        }
4681    }
4682
4683    #[tokio::test]
4684    async fn test_prompt_filter_custom_denial() {
4685        use crate::filter::{CapabilityFilter, DenialBehavior};
4686        use crate::prompt::{Prompt, PromptBuilder};
4687        use std::collections::HashMap;
4688
4689        let admin_prompt = PromptBuilder::new("system_debug")
4690            .description("Admin prompt")
4691            .user_message("Debug");
4692
4693        let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
4694            CapabilityFilter::new(|_, _: &Prompt| false)
4695                .denial_behavior(DenialBehavior::Unauthorized),
4696        );
4697
4698        // Initialize session
4699        init_router(&mut router).await;
4700
4701        let req = RouterRequest {
4702            id: RequestId::Number(1),
4703            inner: McpRequest::GetPrompt(GetPromptParams {
4704                name: "system_debug".to_string(),
4705                arguments: HashMap::new(),
4706                meta: None,
4707            }),
4708            extensions: Extensions::new(),
4709        };
4710
4711        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4712
4713        // Should get forbidden error
4714        match resp.inner {
4715            Err(e) => {
4716                assert_eq!(e.code, -32007); // Forbidden
4717                assert!(e.message.contains("Unauthorized"));
4718            }
4719            _ => panic!("Expected JsonRpc error"),
4720        }
4721    }
4722
4723    // =========================================================================
4724    // Router Composition Tests (merge/nest)
4725    // =========================================================================
4726
4727    #[derive(Debug, Deserialize, JsonSchema)]
4728    struct StringInput {
4729        value: String,
4730    }
4731
4732    #[tokio::test]
4733    async fn test_router_merge_tools() {
4734        // Create first router with a tool
4735        let tool_a = ToolBuilder::new("tool_a")
4736            .description("Tool A")
4737            .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
4738            .build();
4739
4740        let router_a = McpRouter::new().tool(tool_a);
4741
4742        // Create second router with different tools
4743        let tool_b = ToolBuilder::new("tool_b")
4744            .description("Tool B")
4745            .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
4746            .build();
4747        let tool_c = ToolBuilder::new("tool_c")
4748            .description("Tool C")
4749            .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
4750            .build();
4751
4752        let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
4753
4754        // Merge them
4755        let mut merged = McpRouter::new()
4756            .server_info("merged", "1.0")
4757            .merge(router_a)
4758            .merge(router_b);
4759
4760        init_router(&mut merged).await;
4761
4762        // List tools
4763        let req = RouterRequest {
4764            id: RequestId::Number(1),
4765            inner: McpRequest::ListTools(ListToolsParams::default()),
4766            extensions: Extensions::new(),
4767        };
4768
4769        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4770
4771        match resp.inner {
4772            Ok(McpResponse::ListTools(result)) => {
4773                assert_eq!(result.tools.len(), 3);
4774                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4775                assert!(names.contains(&"tool_a"));
4776                assert!(names.contains(&"tool_b"));
4777                assert!(names.contains(&"tool_c"));
4778            }
4779            _ => panic!("Expected ListTools response"),
4780        }
4781    }
4782
4783    #[tokio::test]
4784    async fn test_router_merge_overwrites_duplicates() {
4785        // Create first router with a tool
4786        let tool_v1 = ToolBuilder::new("shared")
4787            .description("Version 1")
4788            .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
4789            .build();
4790
4791        let router_a = McpRouter::new().tool(tool_v1);
4792
4793        // Create second router with same tool name but different description
4794        let tool_v2 = ToolBuilder::new("shared")
4795            .description("Version 2")
4796            .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
4797            .build();
4798
4799        let router_b = McpRouter::new().tool(tool_v2);
4800
4801        // Merge - second should win
4802        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4803
4804        init_router(&mut merged).await;
4805
4806        let req = RouterRequest {
4807            id: RequestId::Number(1),
4808            inner: McpRequest::ListTools(ListToolsParams::default()),
4809            extensions: Extensions::new(),
4810        };
4811
4812        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4813
4814        match resp.inner {
4815            Ok(McpResponse::ListTools(result)) => {
4816                assert_eq!(result.tools.len(), 1);
4817                assert_eq!(result.tools[0].name, "shared");
4818                assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
4819            }
4820            _ => panic!("Expected ListTools response"),
4821        }
4822    }
4823
4824    #[tokio::test]
4825    async fn test_router_merge_resources() {
4826        use crate::resource::ResourceBuilder;
4827
4828        // Create routers with different resources
4829        let router_a = McpRouter::new().resource(
4830            ResourceBuilder::new("file:///a.txt")
4831                .name("File A")
4832                .text("content a"),
4833        );
4834
4835        let router_b = McpRouter::new().resource(
4836            ResourceBuilder::new("file:///b.txt")
4837                .name("File B")
4838                .text("content b"),
4839        );
4840
4841        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4842
4843        init_router(&mut merged).await;
4844
4845        let req = RouterRequest {
4846            id: RequestId::Number(1),
4847            inner: McpRequest::ListResources(ListResourcesParams::default()),
4848            extensions: Extensions::new(),
4849        };
4850
4851        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4852
4853        match resp.inner {
4854            Ok(McpResponse::ListResources(result)) => {
4855                assert_eq!(result.resources.len(), 2);
4856                let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4857                assert!(uris.contains(&"file:///a.txt"));
4858                assert!(uris.contains(&"file:///b.txt"));
4859            }
4860            _ => panic!("Expected ListResources response"),
4861        }
4862    }
4863
4864    #[tokio::test]
4865    async fn test_router_merge_prompts() {
4866        use crate::prompt::PromptBuilder;
4867
4868        let router_a =
4869            McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4870
4871        let router_b =
4872            McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4873
4874        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4875
4876        init_router(&mut merged).await;
4877
4878        let req = RouterRequest {
4879            id: RequestId::Number(1),
4880            inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4881            extensions: Extensions::new(),
4882        };
4883
4884        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4885
4886        match resp.inner {
4887            Ok(McpResponse::ListPrompts(result)) => {
4888                assert_eq!(result.prompts.len(), 2);
4889                let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4890                assert!(names.contains(&"prompt_a"));
4891                assert!(names.contains(&"prompt_b"));
4892            }
4893            _ => panic!("Expected ListPrompts response"),
4894        }
4895    }
4896
4897    #[tokio::test]
4898    async fn test_router_nest_prefixes_tools() {
4899        // Create a router with tools
4900        let tool_query = ToolBuilder::new("query")
4901            .description("Query the database")
4902            .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4903            .build();
4904        let tool_insert = ToolBuilder::new("insert")
4905            .description("Insert into database")
4906            .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4907            .build();
4908
4909        let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4910
4911        // Nest under "db" prefix
4912        let mut router = McpRouter::new()
4913            .server_info("nested", "1.0")
4914            .nest("db", db_router);
4915
4916        init_router(&mut router).await;
4917
4918        let req = RouterRequest {
4919            id: RequestId::Number(1),
4920            inner: McpRequest::ListTools(ListToolsParams::default()),
4921            extensions: Extensions::new(),
4922        };
4923
4924        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4925
4926        match resp.inner {
4927            Ok(McpResponse::ListTools(result)) => {
4928                assert_eq!(result.tools.len(), 2);
4929                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4930                assert!(names.contains(&"db.query"));
4931                assert!(names.contains(&"db.insert"));
4932            }
4933            _ => panic!("Expected ListTools response"),
4934        }
4935    }
4936
4937    #[tokio::test]
4938    async fn test_router_nest_call_prefixed_tool() {
4939        let tool = ToolBuilder::new("echo")
4940            .description("Echo input")
4941            .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4942            .build();
4943
4944        let nested_router = McpRouter::new().tool(tool);
4945
4946        let mut router = McpRouter::new().nest("api", nested_router);
4947
4948        init_router(&mut router).await;
4949
4950        // Call the prefixed tool
4951        let req = RouterRequest {
4952            id: RequestId::Number(1),
4953            inner: McpRequest::CallTool(CallToolParams {
4954                name: "api.echo".to_string(),
4955                arguments: serde_json::json!({"value": "hello world"}),
4956                meta: None,
4957                task: None,
4958            }),
4959            extensions: Extensions::new(),
4960        };
4961
4962        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4963
4964        match resp.inner {
4965            Ok(McpResponse::CallTool(result)) => {
4966                assert!(!result.is_error);
4967                match &result.content[0] {
4968                    Content::Text { text, .. } => assert_eq!(text, "hello world"),
4969                    _ => panic!("Expected text content"),
4970                }
4971            }
4972            _ => panic!("Expected CallTool response"),
4973        }
4974    }
4975
4976    #[tokio::test]
4977    async fn test_router_multiple_nests() {
4978        let db_tool = ToolBuilder::new("query")
4979            .description("Database query")
4980            .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4981            .build();
4982
4983        let api_tool = ToolBuilder::new("fetch")
4984            .description("API fetch")
4985            .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4986            .build();
4987
4988        let db_router = McpRouter::new().tool(db_tool);
4989        let api_router = McpRouter::new().tool(api_tool);
4990
4991        let mut router = McpRouter::new()
4992            .nest("db", db_router)
4993            .nest("api", api_router);
4994
4995        init_router(&mut router).await;
4996
4997        let req = RouterRequest {
4998            id: RequestId::Number(1),
4999            inner: McpRequest::ListTools(ListToolsParams::default()),
5000            extensions: Extensions::new(),
5001        };
5002
5003        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5004
5005        match resp.inner {
5006            Ok(McpResponse::ListTools(result)) => {
5007                assert_eq!(result.tools.len(), 2);
5008                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5009                assert!(names.contains(&"db.query"));
5010                assert!(names.contains(&"api.fetch"));
5011            }
5012            _ => panic!("Expected ListTools response"),
5013        }
5014    }
5015
5016    #[tokio::test]
5017    async fn test_router_merge_and_nest_combined() {
5018        // Test combining merge and nest
5019        let tool_a = ToolBuilder::new("local")
5020            .description("Local tool")
5021            .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
5022            .build();
5023
5024        let nested_tool = ToolBuilder::new("remote")
5025            .description("Remote tool")
5026            .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
5027            .build();
5028
5029        let nested_router = McpRouter::new().tool(nested_tool);
5030
5031        let mut router = McpRouter::new()
5032            .tool(tool_a)
5033            .nest("external", nested_router);
5034
5035        init_router(&mut router).await;
5036
5037        let req = RouterRequest {
5038            id: RequestId::Number(1),
5039            inner: McpRequest::ListTools(ListToolsParams::default()),
5040            extensions: Extensions::new(),
5041        };
5042
5043        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5044
5045        match resp.inner {
5046            Ok(McpResponse::ListTools(result)) => {
5047                assert_eq!(result.tools.len(), 2);
5048                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5049                assert!(names.contains(&"local"));
5050                assert!(names.contains(&"external.remote"));
5051            }
5052            _ => panic!("Expected ListTools response"),
5053        }
5054    }
5055
5056    #[tokio::test]
5057    async fn test_router_merge_preserves_server_info() {
5058        let child_router = McpRouter::new()
5059            .server_info("child", "2.0")
5060            .instructions("Child instructions");
5061
5062        let mut router = McpRouter::new()
5063            .server_info("parent", "1.0")
5064            .instructions("Parent instructions")
5065            .merge(child_router);
5066
5067        init_router(&mut router).await;
5068
5069        // Initialize response should have parent's server info
5070        let init_req = RouterRequest {
5071            id: RequestId::Number(99),
5072            inner: McpRequest::Initialize(InitializeParams {
5073                protocol_version: "2025-11-25".to_string(),
5074                capabilities: ClientCapabilities::default(),
5075                client_info: Implementation {
5076                    name: "test".to_string(),
5077                    version: "1.0".to_string(),
5078                    ..Default::default()
5079                },
5080                meta: None,
5081            }),
5082            extensions: Extensions::new(),
5083        };
5084
5085        // Create fresh router for this test since we need to call initialize
5086        let child_router2 = McpRouter::new().server_info("child", "2.0");
5087        let mut fresh_router = McpRouter::new()
5088            .server_info("parent", "1.0")
5089            .merge(child_router2);
5090
5091        let resp = fresh_router
5092            .ready()
5093            .await
5094            .unwrap()
5095            .call(init_req)
5096            .await
5097            .unwrap();
5098
5099        match resp.inner {
5100            Ok(McpResponse::Initialize(result)) => {
5101                assert_eq!(result.server_info.name, "parent");
5102                assert_eq!(result.server_info.version, "1.0");
5103            }
5104            _ => panic!("Expected Initialize response"),
5105        }
5106    }
5107
5108    // =========================================================================
5109    // Auto-instructions tests
5110    // =========================================================================
5111
5112    #[tokio::test]
5113    async fn test_auto_instructions_tools_only() {
5114        let tool_a = ToolBuilder::new("alpha")
5115            .description("Alpha tool")
5116            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5117            .build();
5118        let tool_b = ToolBuilder::new("beta")
5119            .description("Beta tool")
5120            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5121            .build();
5122
5123        let mut router = McpRouter::new()
5124            .auto_instructions()
5125            .tool(tool_a)
5126            .tool(tool_b);
5127
5128        let resp = send_initialize(&mut router).await;
5129        let instructions = resp.instructions.expect("should have instructions");
5130
5131        assert!(instructions.contains("## Tools"));
5132        assert!(instructions.contains("- **alpha**: Alpha tool"));
5133        assert!(instructions.contains("- **beta**: Beta tool"));
5134        // No resources or prompts sections
5135        assert!(!instructions.contains("## Resources"));
5136        assert!(!instructions.contains("## Prompts"));
5137    }
5138
5139    #[tokio::test]
5140    async fn test_auto_instructions_with_annotations() {
5141        let read_only_tool = ToolBuilder::new("query")
5142            .description("Run a query")
5143            .read_only()
5144            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5145            .build();
5146        let destructive_tool = ToolBuilder::new("delete")
5147            .description("Delete a record")
5148            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5149            .build();
5150        let idempotent_tool = ToolBuilder::new("upsert")
5151            .description("Upsert a record")
5152            .non_destructive()
5153            .idempotent()
5154            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5155            .build();
5156
5157        let mut router = McpRouter::new()
5158            .auto_instructions()
5159            .tool(read_only_tool)
5160            .tool(destructive_tool)
5161            .tool(idempotent_tool);
5162
5163        let resp = send_initialize(&mut router).await;
5164        let instructions = resp.instructions.unwrap();
5165
5166        assert!(instructions.contains("- **query**: Run a query [read-only]"));
5167        // delete has no annotations set via builder, so no tags
5168        assert!(instructions.contains("- **delete**: Delete a record\n"));
5169        assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
5170    }
5171
5172    #[tokio::test]
5173    async fn test_auto_instructions_with_resources() {
5174        use crate::resource::ResourceBuilder;
5175
5176        let resource = ResourceBuilder::new("file:///schema.sql")
5177            .name("Schema")
5178            .description("Database schema")
5179            .text("CREATE TABLE ...");
5180
5181        let mut router = McpRouter::new().auto_instructions().resource(resource);
5182
5183        let resp = send_initialize(&mut router).await;
5184        let instructions = resp.instructions.unwrap();
5185
5186        assert!(instructions.contains("## Resources"));
5187        assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
5188        assert!(!instructions.contains("## Tools"));
5189    }
5190
5191    #[tokio::test]
5192    async fn test_auto_instructions_with_resource_templates() {
5193        use crate::resource::ResourceTemplateBuilder;
5194
5195        let template = ResourceTemplateBuilder::new("file:///{path}")
5196            .name("File")
5197            .description("Read a file by path")
5198            .handler(
5199                |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
5200                    Ok(crate::ReadResourceResult::text("content", "text/plain"))
5201                },
5202            );
5203
5204        let mut router = McpRouter::new()
5205            .auto_instructions()
5206            .resource_template(template);
5207
5208        let resp = send_initialize(&mut router).await;
5209        let instructions = resp.instructions.unwrap();
5210
5211        assert!(instructions.contains("## Resources"));
5212        assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
5213    }
5214
5215    #[tokio::test]
5216    async fn test_auto_instructions_with_prompts() {
5217        use crate::prompt::PromptBuilder;
5218
5219        let prompt = PromptBuilder::new("write_query")
5220            .description("Help write a SQL query")
5221            .user_message("Write a query for: {task}");
5222
5223        let mut router = McpRouter::new().auto_instructions().prompt(prompt);
5224
5225        let resp = send_initialize(&mut router).await;
5226        let instructions = resp.instructions.unwrap();
5227
5228        assert!(instructions.contains("## Prompts"));
5229        assert!(instructions.contains("- **write_query**: Help write a SQL query"));
5230        assert!(!instructions.contains("## Tools"));
5231    }
5232
5233    #[tokio::test]
5234    async fn test_auto_instructions_all_sections() {
5235        use crate::prompt::PromptBuilder;
5236        use crate::resource::ResourceBuilder;
5237
5238        let tool = ToolBuilder::new("query")
5239            .description("Execute SQL")
5240            .read_only()
5241            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5242            .build();
5243        let resource = ResourceBuilder::new("db://schema")
5244            .name("Schema")
5245            .description("Full database schema")
5246            .text("schema");
5247        let prompt = PromptBuilder::new("write_query")
5248            .description("Help write a SQL query")
5249            .user_message("Write a query");
5250
5251        let mut router = McpRouter::new()
5252            .auto_instructions()
5253            .tool(tool)
5254            .resource(resource)
5255            .prompt(prompt);
5256
5257        let resp = send_initialize(&mut router).await;
5258        let instructions = resp.instructions.unwrap();
5259
5260        // All three sections present
5261        assert!(instructions.contains("## Tools"));
5262        assert!(instructions.contains("## Resources"));
5263        assert!(instructions.contains("## Prompts"));
5264
5265        // Sections appear in order: Tools, Resources, Prompts
5266        let tools_pos = instructions.find("## Tools").unwrap();
5267        let resources_pos = instructions.find("## Resources").unwrap();
5268        let prompts_pos = instructions.find("## Prompts").unwrap();
5269        assert!(tools_pos < resources_pos);
5270        assert!(resources_pos < prompts_pos);
5271    }
5272
5273    #[tokio::test]
5274    async fn test_auto_instructions_with_prefix_and_suffix() {
5275        let tool = ToolBuilder::new("echo")
5276            .description("Echo input")
5277            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5278            .build();
5279
5280        let mut router = McpRouter::new()
5281            .auto_instructions_with(
5282                Some("This server provides echo capabilities."),
5283                Some("Contact admin@example.com for support."),
5284            )
5285            .tool(tool);
5286
5287        let resp = send_initialize(&mut router).await;
5288        let instructions = resp.instructions.unwrap();
5289
5290        assert!(instructions.starts_with("This server provides echo capabilities."));
5291        assert!(instructions.ends_with("Contact admin@example.com for support."));
5292        assert!(instructions.contains("## Tools"));
5293        assert!(instructions.contains("- **echo**: Echo input"));
5294    }
5295
5296    #[tokio::test]
5297    async fn test_auto_instructions_prefix_only() {
5298        let tool = ToolBuilder::new("echo")
5299            .description("Echo input")
5300            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5301            .build();
5302
5303        let mut router = McpRouter::new()
5304            .auto_instructions_with(Some("My server intro."), None::<String>)
5305            .tool(tool);
5306
5307        let resp = send_initialize(&mut router).await;
5308        let instructions = resp.instructions.unwrap();
5309
5310        assert!(instructions.starts_with("My server intro."));
5311        assert!(instructions.contains("- **echo**: Echo input"));
5312    }
5313
5314    #[tokio::test]
5315    async fn test_auto_instructions_empty_router() {
5316        let mut router = McpRouter::new().auto_instructions();
5317
5318        let resp = send_initialize(&mut router).await;
5319        let instructions = resp.instructions.expect("should have instructions");
5320
5321        // No sections when nothing is registered
5322        assert!(!instructions.contains("## Tools"));
5323        assert!(!instructions.contains("## Resources"));
5324        assert!(!instructions.contains("## Prompts"));
5325        assert!(instructions.is_empty());
5326    }
5327
5328    #[tokio::test]
5329    async fn test_auto_instructions_overrides_manual() {
5330        let tool = ToolBuilder::new("echo")
5331            .description("Echo input")
5332            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5333            .build();
5334
5335        let mut router = McpRouter::new()
5336            .instructions("This will be overridden")
5337            .auto_instructions()
5338            .tool(tool);
5339
5340        let resp = send_initialize(&mut router).await;
5341        let instructions = resp.instructions.unwrap();
5342
5343        assert!(!instructions.contains("This will be overridden"));
5344        assert!(instructions.contains("- **echo**: Echo input"));
5345    }
5346
5347    #[tokio::test]
5348    async fn test_no_auto_instructions_returns_manual() {
5349        let tool = ToolBuilder::new("echo")
5350            .description("Echo input")
5351            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5352            .build();
5353
5354        let mut router = McpRouter::new()
5355            .instructions("Manual instructions here")
5356            .tool(tool);
5357
5358        let resp = send_initialize(&mut router).await;
5359        let instructions = resp.instructions.unwrap();
5360
5361        assert_eq!(instructions, "Manual instructions here");
5362    }
5363
5364    #[tokio::test]
5365    async fn test_auto_instructions_no_description_fallback() {
5366        let tool = ToolBuilder::new("mystery")
5367            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5368            .build();
5369
5370        let mut router = McpRouter::new().auto_instructions().tool(tool);
5371
5372        let resp = send_initialize(&mut router).await;
5373        let instructions = resp.instructions.unwrap();
5374
5375        assert!(instructions.contains("- **mystery**: No description"));
5376    }
5377
5378    #[tokio::test]
5379    async fn test_auto_instructions_sorted_alphabetically() {
5380        let tool_z = ToolBuilder::new("zebra")
5381            .description("Z tool")
5382            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5383            .build();
5384        let tool_a = ToolBuilder::new("alpha")
5385            .description("A tool")
5386            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5387            .build();
5388        let tool_m = ToolBuilder::new("middle")
5389            .description("M tool")
5390            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5391            .build();
5392
5393        let mut router = McpRouter::new()
5394            .auto_instructions()
5395            .tool(tool_z)
5396            .tool(tool_a)
5397            .tool(tool_m);
5398
5399        let resp = send_initialize(&mut router).await;
5400        let instructions = resp.instructions.unwrap();
5401
5402        let alpha_pos = instructions.find("**alpha**").unwrap();
5403        let middle_pos = instructions.find("**middle**").unwrap();
5404        let zebra_pos = instructions.find("**zebra**").unwrap();
5405        assert!(alpha_pos < middle_pos);
5406        assert!(middle_pos < zebra_pos);
5407    }
5408
5409    #[tokio::test]
5410    async fn test_auto_instructions_read_only_and_idempotent_tags() {
5411        let tool = ToolBuilder::new("safe_update")
5412            .description("Safe update operation")
5413            .idempotent()
5414            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5415            .build();
5416
5417        let mut router = McpRouter::new().auto_instructions().tool(tool);
5418
5419        let resp = send_initialize(&mut router).await;
5420        let instructions = resp.instructions.unwrap();
5421
5422        assert!(
5423            instructions.contains("[idempotent]"),
5424            "got: {}",
5425            instructions
5426        );
5427    }
5428
5429    #[tokio::test]
5430    async fn test_auto_instructions_lazy_generation() {
5431        // auto_instructions() is called BEFORE tools are registered
5432        // but instructions should still include tools
5433        let mut router = McpRouter::new().auto_instructions();
5434
5435        let tool = ToolBuilder::new("late_tool")
5436            .description("Added after auto_instructions")
5437            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5438            .build();
5439
5440        router = router.tool(tool);
5441
5442        let resp = send_initialize(&mut router).await;
5443        let instructions = resp.instructions.unwrap();
5444
5445        assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
5446    }
5447
5448    #[tokio::test]
5449    async fn test_auto_instructions_multiple_annotation_tags() {
5450        let tool = ToolBuilder::new("update")
5451            .description("Update a record")
5452            .annotations(ToolAnnotations {
5453                read_only_hint: true,
5454                idempotent_hint: true,
5455                ..Default::default()
5456            })
5457            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5458            .build();
5459
5460        let mut router = McpRouter::new().auto_instructions().tool(tool);
5461
5462        let resp = send_initialize(&mut router).await;
5463        let instructions = resp.instructions.unwrap();
5464
5465        assert!(
5466            instructions.contains("[read-only, idempotent]"),
5467            "got: {}",
5468            instructions
5469        );
5470    }
5471
5472    #[tokio::test]
5473    async fn test_auto_instructions_no_annotations_no_tags() {
5474        // Tools without annotations should have no tags at all
5475        let tool = ToolBuilder::new("fetch")
5476            .description("Fetch data")
5477            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5478            .build();
5479
5480        let mut router = McpRouter::new().auto_instructions().tool(tool);
5481
5482        let resp = send_initialize(&mut router).await;
5483        let instructions = resp.instructions.unwrap();
5484
5485        // No bracket tags
5486        assert!(
5487            !instructions.contains('['),
5488            "should have no tags, got: {}",
5489            instructions
5490        );
5491        assert!(instructions.contains("- **fetch**: Fetch data"));
5492    }
5493
5494    /// Helper to send an Initialize request and return the result
5495    async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
5496        let init_req = RouterRequest {
5497            id: RequestId::Number(0),
5498            inner: McpRequest::Initialize(InitializeParams {
5499                protocol_version: "2025-11-25".to_string(),
5500                capabilities: ClientCapabilities {
5501                    roots: None,
5502                    sampling: None,
5503                    elicitation: None,
5504                    tasks: None,
5505                    experimental: None,
5506                    extensions: None,
5507                },
5508                client_info: Implementation {
5509                    name: "test".to_string(),
5510                    version: "1.0".to_string(),
5511                    ..Default::default()
5512                },
5513                meta: None,
5514            }),
5515            extensions: Extensions::new(),
5516        };
5517        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5518        match resp.inner {
5519            Ok(McpResponse::Initialize(result)) => result,
5520            other => panic!("Expected Initialize response, got {:?}", other),
5521        }
5522    }
5523
5524    #[tokio::test]
5525    async fn test_notify_tools_list_changed() {
5526        let (tx, mut rx) = crate::context::notification_channel(16);
5527
5528        let router = McpRouter::new()
5529            .server_info("test", "1.0")
5530            .with_notification_sender(tx);
5531
5532        assert!(router.notify_tools_list_changed());
5533
5534        let notification = rx.recv().await.unwrap();
5535        assert!(matches!(notification, ServerNotification::ToolsListChanged));
5536    }
5537
5538    #[tokio::test]
5539    async fn test_notify_prompts_list_changed() {
5540        let (tx, mut rx) = crate::context::notification_channel(16);
5541
5542        let router = McpRouter::new()
5543            .server_info("test", "1.0")
5544            .with_notification_sender(tx);
5545
5546        assert!(router.notify_prompts_list_changed());
5547
5548        let notification = rx.recv().await.unwrap();
5549        assert!(matches!(
5550            notification,
5551            ServerNotification::PromptsListChanged
5552        ));
5553    }
5554
5555    #[tokio::test]
5556    async fn test_notify_without_sender_returns_false() {
5557        let router = McpRouter::new().server_info("test", "1.0");
5558
5559        assert!(!router.notify_tools_list_changed());
5560        assert!(!router.notify_prompts_list_changed());
5561        assert!(!router.notify_resources_list_changed());
5562    }
5563
5564    #[tokio::test]
5565    async fn test_list_changed_capabilities_with_notification_sender() {
5566        let (tx, _rx) = crate::context::notification_channel(16);
5567        let tool = ToolBuilder::new("test")
5568            .description("test")
5569            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5570            .build();
5571
5572        let mut router = McpRouter::new()
5573            .server_info("test", "1.0")
5574            .tool(tool)
5575            .with_notification_sender(tx);
5576
5577        init_router(&mut router).await;
5578
5579        let caps = router.capabilities();
5580        let tools_cap = caps.tools.expect("tools capability should be present");
5581        assert!(
5582            tools_cap.list_changed,
5583            "tools.listChanged should be true when notification sender is configured"
5584        );
5585    }
5586
5587    #[tokio::test]
5588    async fn test_list_changed_capabilities_without_notification_sender() {
5589        let tool = ToolBuilder::new("test")
5590            .description("test")
5591            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5592            .build();
5593
5594        let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
5595
5596        init_router(&mut router).await;
5597
5598        let caps = router.capabilities();
5599        let tools_cap = caps.tools.expect("tools capability should be present");
5600        assert!(
5601            !tools_cap.list_changed,
5602            "tools.listChanged should be false without notification sender"
5603        );
5604    }
5605
5606    #[tokio::test]
5607    async fn test_set_logging_level_filters_messages() {
5608        let (tx, mut rx) = crate::context::notification_channel(16);
5609
5610        let mut router = McpRouter::new()
5611            .server_info("test", "1.0")
5612            .with_notification_sender(tx);
5613
5614        init_router(&mut router).await;
5615
5616        // Set logging level to Warning
5617        let set_level_req = RouterRequest {
5618            id: RequestId::Number(99),
5619            inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
5620                level: LogLevel::Warning,
5621                meta: None,
5622            }),
5623            extensions: crate::context::Extensions::new(),
5624        };
5625        let resp = router
5626            .ready()
5627            .await
5628            .unwrap()
5629            .call(set_level_req)
5630            .await
5631            .unwrap();
5632        assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
5633
5634        // Create a context from the router (simulating a handler)
5635        let ctx = router.create_context(RequestId::Number(100), None);
5636
5637        // Error (more severe than Warning) should pass through
5638        ctx.send_log(LoggingMessageParams::new(
5639            LogLevel::Error,
5640            serde_json::Value::Null,
5641        ));
5642        assert!(
5643            rx.try_recv().is_ok(),
5644            "Error should pass through Warning filter"
5645        );
5646
5647        // Info (less severe than Warning) should be filtered
5648        ctx.send_log(LoggingMessageParams::new(
5649            LogLevel::Info,
5650            serde_json::Value::Null,
5651        ));
5652        assert!(
5653            rx.try_recv().is_err(),
5654            "Info should be filtered at Warning level"
5655        );
5656    }
5657
5658    #[test]
5659    fn test_paginate_no_page_size() {
5660        let items = vec![1, 2, 3, 4, 5];
5661        let (page, cursor) = paginate(items.clone(), None, None).unwrap();
5662        assert_eq!(page, items);
5663        assert!(cursor.is_none());
5664    }
5665
5666    #[test]
5667    fn test_paginate_first_page() {
5668        let items = vec![1, 2, 3, 4, 5];
5669        let (page, cursor) = paginate(items, None, Some(2)).unwrap();
5670        assert_eq!(page, vec![1, 2]);
5671        assert!(cursor.is_some());
5672    }
5673
5674    #[test]
5675    fn test_paginate_middle_page() {
5676        let items = vec![1, 2, 3, 4, 5];
5677        let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
5678        assert_eq!(page1, vec![1, 2]);
5679
5680        let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
5681        assert_eq!(page2, vec![3, 4]);
5682        assert!(cursor2.is_some());
5683    }
5684
5685    #[test]
5686    fn test_paginate_last_page() {
5687        let items = vec![1, 2, 3, 4, 5];
5688        // Skip to offset 4 (last item)
5689        let cursor = encode_cursor(4);
5690        let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
5691        assert_eq!(page, vec![5]);
5692        assert!(next.is_none());
5693    }
5694
5695    #[test]
5696    fn test_paginate_exact_boundary() {
5697        let items = vec![1, 2, 3, 4];
5698        let (page, cursor) = paginate(items, None, Some(4)).unwrap();
5699        assert_eq!(page, vec![1, 2, 3, 4]);
5700        assert!(cursor.is_none());
5701    }
5702
5703    #[test]
5704    fn test_paginate_invalid_cursor() {
5705        let items = vec![1, 2, 3];
5706        let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
5707        assert!(result.is_err());
5708    }
5709
5710    #[test]
5711    fn test_cursor_round_trip() {
5712        let offset = 42;
5713        let encoded = encode_cursor(offset);
5714        let decoded = decode_cursor(&encoded).unwrap();
5715        assert_eq!(decoded, offset);
5716    }
5717
5718    #[tokio::test]
5719    async fn test_list_tools_pagination() {
5720        let tool_a = ToolBuilder::new("alpha")
5721            .description("a")
5722            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5723            .build();
5724        let tool_b = ToolBuilder::new("beta")
5725            .description("b")
5726            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5727            .build();
5728        let tool_c = ToolBuilder::new("gamma")
5729            .description("c")
5730            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5731            .build();
5732
5733        let mut router = McpRouter::new()
5734            .server_info("test", "1.0")
5735            .page_size(2)
5736            .tool(tool_a)
5737            .tool(tool_b)
5738            .tool(tool_c);
5739
5740        init_router(&mut router).await;
5741
5742        // First page
5743        let req = RouterRequest {
5744            id: RequestId::Number(1),
5745            inner: McpRequest::ListTools(ListToolsParams {
5746                cursor: None,
5747                meta: None,
5748            }),
5749            extensions: Extensions::new(),
5750        };
5751        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5752        let (tools, next_cursor) = match resp.inner {
5753            Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5754            other => panic!("Expected ListTools, got {:?}", other),
5755        };
5756        assert_eq!(tools.len(), 2);
5757        assert_eq!(tools[0].name, "alpha");
5758        assert_eq!(tools[1].name, "beta");
5759        assert!(next_cursor.is_some());
5760
5761        // Second page
5762        let req = RouterRequest {
5763            id: RequestId::Number(2),
5764            inner: McpRequest::ListTools(ListToolsParams {
5765                cursor: next_cursor,
5766                meta: None,
5767            }),
5768            extensions: Extensions::new(),
5769        };
5770        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5771        let (tools, next_cursor) = match resp.inner {
5772            Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5773            other => panic!("Expected ListTools, got {:?}", other),
5774        };
5775        assert_eq!(tools.len(), 1);
5776        assert_eq!(tools[0].name, "gamma");
5777        assert!(next_cursor.is_none());
5778    }
5779
5780    #[tokio::test]
5781    async fn test_list_tools_no_pagination_by_default() {
5782        let tool_a = ToolBuilder::new("alpha")
5783            .description("a")
5784            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5785            .build();
5786        let tool_b = ToolBuilder::new("beta")
5787            .description("b")
5788            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5789            .build();
5790
5791        let mut router = McpRouter::new()
5792            .server_info("test", "1.0")
5793            .tool(tool_a)
5794            .tool(tool_b);
5795
5796        init_router(&mut router).await;
5797
5798        let req = RouterRequest {
5799            id: RequestId::Number(1),
5800            inner: McpRequest::ListTools(ListToolsParams {
5801                cursor: None,
5802                meta: None,
5803            }),
5804            extensions: Extensions::new(),
5805        };
5806        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5807        match resp.inner {
5808            Ok(McpResponse::ListTools(result)) => {
5809                assert_eq!(result.tools.len(), 2);
5810                assert!(result.next_cursor.is_none());
5811            }
5812            other => panic!("Expected ListTools, got {:?}", other),
5813        }
5814    }
5815
5816    // =========================================================================
5817    // Dynamic Tool Registry Tests
5818    // =========================================================================
5819
5820    #[cfg(feature = "dynamic-tools")]
5821    mod dynamic_tools_tests {
5822        use super::*;
5823
5824        #[tokio::test]
5825        async fn test_dynamic_tools_register_and_list() {
5826            let (router, registry) = McpRouter::new()
5827                .server_info("test", "1.0")
5828                .with_dynamic_tools();
5829
5830            let tool = ToolBuilder::new("dynamic_echo")
5831                .description("Dynamic echo")
5832                .handler(|input: AddInput| async move {
5833                    Ok(CallToolResult::text(format!("{}", input.a)))
5834                })
5835                .build();
5836
5837            registry.register(tool);
5838
5839            let mut router = router;
5840            init_router(&mut router).await;
5841
5842            let req = RouterRequest {
5843                id: RequestId::Number(1),
5844                inner: McpRequest::ListTools(ListToolsParams::default()),
5845                extensions: Extensions::new(),
5846            };
5847
5848            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5849            match resp.inner {
5850                Ok(McpResponse::ListTools(result)) => {
5851                    assert_eq!(result.tools.len(), 1);
5852                    assert_eq!(result.tools[0].name, "dynamic_echo");
5853                }
5854                _ => panic!("Expected ListTools response"),
5855            }
5856        }
5857
5858        #[tokio::test]
5859        async fn test_dynamic_tools_unregister() {
5860            let (router, registry) = McpRouter::new()
5861                .server_info("test", "1.0")
5862                .with_dynamic_tools();
5863
5864            let tool = ToolBuilder::new("temp")
5865                .description("Temporary")
5866                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5867                .build();
5868
5869            registry.register(tool);
5870            assert!(registry.contains("temp"));
5871
5872            let removed = registry.unregister("temp");
5873            assert!(removed);
5874            assert!(!registry.contains("temp"));
5875
5876            // Unregistering again returns false
5877            assert!(!registry.unregister("temp"));
5878
5879            let mut router = router;
5880            init_router(&mut router).await;
5881
5882            let req = RouterRequest {
5883                id: RequestId::Number(1),
5884                inner: McpRequest::ListTools(ListToolsParams::default()),
5885                extensions: Extensions::new(),
5886            };
5887
5888            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5889            match resp.inner {
5890                Ok(McpResponse::ListTools(result)) => {
5891                    assert_eq!(result.tools.len(), 0);
5892                }
5893                _ => panic!("Expected ListTools response"),
5894            }
5895        }
5896
5897        #[tokio::test]
5898        async fn test_dynamic_tools_merged_with_static() {
5899            let static_tool = ToolBuilder::new("static_tool")
5900                .description("Static")
5901                .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5902                .build();
5903
5904            let (router, registry) = McpRouter::new()
5905                .server_info("test", "1.0")
5906                .tool(static_tool)
5907                .with_dynamic_tools();
5908
5909            let dynamic_tool = ToolBuilder::new("dynamic_tool")
5910                .description("Dynamic")
5911                .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5912                .build();
5913
5914            registry.register(dynamic_tool);
5915
5916            let mut router = router;
5917            init_router(&mut router).await;
5918
5919            let req = RouterRequest {
5920                id: RequestId::Number(1),
5921                inner: McpRequest::ListTools(ListToolsParams::default()),
5922                extensions: Extensions::new(),
5923            };
5924
5925            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5926            match resp.inner {
5927                Ok(McpResponse::ListTools(result)) => {
5928                    assert_eq!(result.tools.len(), 2);
5929                    let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5930                    assert!(names.contains(&"static_tool"));
5931                    assert!(names.contains(&"dynamic_tool"));
5932                }
5933                _ => panic!("Expected ListTools response"),
5934            }
5935        }
5936
5937        #[tokio::test]
5938        async fn test_static_tools_shadow_dynamic() {
5939            let static_tool = ToolBuilder::new("shared")
5940                .description("Static version")
5941                .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5942                .build();
5943
5944            let (router, registry) = McpRouter::new()
5945                .server_info("test", "1.0")
5946                .tool(static_tool)
5947                .with_dynamic_tools();
5948
5949            let dynamic_tool = ToolBuilder::new("shared")
5950                .description("Dynamic version")
5951                .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5952                .build();
5953
5954            registry.register(dynamic_tool);
5955
5956            let mut router = router;
5957            init_router(&mut router).await;
5958
5959            // List should only show the static version
5960            let req = RouterRequest {
5961                id: RequestId::Number(1),
5962                inner: McpRequest::ListTools(ListToolsParams::default()),
5963                extensions: Extensions::new(),
5964            };
5965
5966            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5967            match resp.inner {
5968                Ok(McpResponse::ListTools(result)) => {
5969                    assert_eq!(result.tools.len(), 1);
5970                    assert_eq!(result.tools[0].name, "shared");
5971                    assert_eq!(
5972                        result.tools[0].description.as_deref(),
5973                        Some("Static version")
5974                    );
5975                }
5976                _ => panic!("Expected ListTools response"),
5977            }
5978
5979            // Call should dispatch to the static tool
5980            let req = RouterRequest {
5981                id: RequestId::Number(2),
5982                inner: McpRequest::CallTool(CallToolParams {
5983                    name: "shared".to_string(),
5984                    arguments: serde_json::json!({"a": 1, "b": 2}),
5985                    meta: None,
5986                    task: None,
5987                }),
5988                extensions: Extensions::new(),
5989            };
5990
5991            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5992            match resp.inner {
5993                Ok(McpResponse::CallTool(result)) => {
5994                    assert!(!result.is_error);
5995                    match &result.content[0] {
5996                        Content::Text { text, .. } => assert_eq!(text, "static"),
5997                        _ => panic!("Expected text content"),
5998                    }
5999                }
6000                _ => panic!("Expected CallTool response"),
6001            }
6002        }
6003
6004        #[tokio::test]
6005        async fn test_dynamic_tools_call() {
6006            let (router, registry) = McpRouter::new()
6007                .server_info("test", "1.0")
6008                .with_dynamic_tools();
6009
6010            let tool = ToolBuilder::new("add")
6011                .description("Add two numbers")
6012                .handler(|input: AddInput| async move {
6013                    Ok(CallToolResult::text(format!("{}", input.a + input.b)))
6014                })
6015                .build();
6016
6017            registry.register(tool);
6018
6019            let mut router = router;
6020            init_router(&mut router).await;
6021
6022            let req = RouterRequest {
6023                id: RequestId::Number(1),
6024                inner: McpRequest::CallTool(CallToolParams {
6025                    name: "add".to_string(),
6026                    arguments: serde_json::json!({"a": 3, "b": 4}),
6027                    meta: None,
6028                    task: None,
6029                }),
6030                extensions: Extensions::new(),
6031            };
6032
6033            let resp = router.ready().await.unwrap().call(req).await.unwrap();
6034            match resp.inner {
6035                Ok(McpResponse::CallTool(result)) => {
6036                    assert!(!result.is_error);
6037                    match &result.content[0] {
6038                        Content::Text { text, .. } => assert_eq!(text, "7"),
6039                        _ => panic!("Expected text content"),
6040                    }
6041                }
6042                _ => panic!("Expected CallTool response"),
6043            }
6044        }
6045
6046        #[tokio::test]
6047        async fn test_dynamic_tools_notification_on_register() {
6048            let (tx, mut rx) = crate::context::notification_channel(16);
6049            let (router, registry) = McpRouter::new()
6050                .server_info("test", "1.0")
6051                .with_dynamic_tools();
6052            let _router = router.with_notification_sender(tx);
6053
6054            let tool = ToolBuilder::new("notified")
6055                .description("Test")
6056                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6057                .build();
6058
6059            registry.register(tool);
6060
6061            let notification = rx.recv().await.unwrap();
6062            assert!(matches!(notification, ServerNotification::ToolsListChanged));
6063        }
6064
6065        #[tokio::test]
6066        async fn test_dynamic_tools_notification_on_unregister() {
6067            let (tx, mut rx) = crate::context::notification_channel(16);
6068            let (router, registry) = McpRouter::new()
6069                .server_info("test", "1.0")
6070                .with_dynamic_tools();
6071            let _router = router.with_notification_sender(tx);
6072
6073            let tool = ToolBuilder::new("notified")
6074                .description("Test")
6075                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6076                .build();
6077
6078            registry.register(tool);
6079            // Consume the register notification
6080            let _ = rx.recv().await.unwrap();
6081
6082            registry.unregister("notified");
6083            let notification = rx.recv().await.unwrap();
6084            assert!(matches!(notification, ServerNotification::ToolsListChanged));
6085        }
6086
6087        #[tokio::test]
6088        async fn test_dynamic_tools_no_notification_on_empty_unregister() {
6089            let (tx, mut rx) = crate::context::notification_channel(16);
6090            let (router, registry) = McpRouter::new()
6091                .server_info("test", "1.0")
6092                .with_dynamic_tools();
6093            let _router = router.with_notification_sender(tx);
6094
6095            // Unregister a tool that doesn't exist — should NOT send notification
6096            assert!(!registry.unregister("nonexistent"));
6097
6098            // Channel should be empty
6099            assert!(rx.try_recv().is_err());
6100        }
6101
6102        #[tokio::test]
6103        async fn test_dynamic_tools_filter_applies() {
6104            use crate::filter::CapabilityFilter;
6105
6106            let (router, registry) = McpRouter::new()
6107                .server_info("test", "1.0")
6108                .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
6109                    tool.name != "hidden"
6110                }))
6111                .with_dynamic_tools();
6112
6113            let visible = ToolBuilder::new("visible")
6114                .description("Visible")
6115                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6116                .build();
6117
6118            let hidden = ToolBuilder::new("hidden")
6119                .description("Hidden")
6120                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6121                .build();
6122
6123            registry.register(visible);
6124            registry.register(hidden);
6125
6126            let mut router = router;
6127            init_router(&mut router).await;
6128
6129            // List should only show visible tool
6130            let req = RouterRequest {
6131                id: RequestId::Number(1),
6132                inner: McpRequest::ListTools(ListToolsParams::default()),
6133                extensions: Extensions::new(),
6134            };
6135
6136            let resp = router.ready().await.unwrap().call(req).await.unwrap();
6137            match resp.inner {
6138                Ok(McpResponse::ListTools(result)) => {
6139                    assert_eq!(result.tools.len(), 1);
6140                    assert_eq!(result.tools[0].name, "visible");
6141                }
6142                _ => panic!("Expected ListTools response"),
6143            }
6144
6145            // Call to hidden tool should be denied
6146            let req = RouterRequest {
6147                id: RequestId::Number(2),
6148                inner: McpRequest::CallTool(CallToolParams {
6149                    name: "hidden".to_string(),
6150                    arguments: serde_json::json!({"a": 1, "b": 2}),
6151                    meta: None,
6152                    task: None,
6153                }),
6154                extensions: Extensions::new(),
6155            };
6156
6157            let resp = router.ready().await.unwrap().call(req).await.unwrap();
6158            match resp.inner {
6159                Err(e) => {
6160                    assert_eq!(e.code, -32601); // Method not found
6161                }
6162                _ => panic!("Expected JsonRpc error"),
6163            }
6164        }
6165
6166        #[tokio::test]
6167        async fn test_dynamic_tools_capabilities_advertised() {
6168            // No static tools, but dynamic tools enabled — should advertise tools capability
6169            let (mut router, _registry) = McpRouter::new()
6170                .server_info("test", "1.0")
6171                .with_dynamic_tools();
6172
6173            let init_req = RouterRequest {
6174                id: RequestId::Number(1),
6175                inner: McpRequest::Initialize(InitializeParams {
6176                    protocol_version: "2025-11-25".to_string(),
6177                    capabilities: ClientCapabilities::default(),
6178                    client_info: Implementation {
6179                        name: "test".to_string(),
6180                        version: "1.0".to_string(),
6181                        ..Default::default()
6182                    },
6183                    meta: None,
6184                }),
6185                extensions: Extensions::new(),
6186            };
6187
6188            let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
6189            match resp.inner {
6190                Ok(McpResponse::Initialize(result)) => {
6191                    assert!(result.capabilities.tools.is_some());
6192                }
6193                _ => panic!("Expected Initialize response"),
6194            }
6195        }
6196
6197        #[tokio::test]
6198        async fn test_dynamic_tools_multi_session_notification() {
6199            let (tx1, mut rx1) = crate::context::notification_channel(16);
6200            let (tx2, mut rx2) = crate::context::notification_channel(16);
6201
6202            let (router, registry) = McpRouter::new()
6203                .server_info("test", "1.0")
6204                .with_dynamic_tools();
6205
6206            // Simulate two sessions by calling with_notification_sender on two clones
6207            let _session1 = router.clone().with_notification_sender(tx1);
6208            let _session2 = router.clone().with_notification_sender(tx2);
6209
6210            let tool = ToolBuilder::new("broadcast")
6211                .description("Test")
6212                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6213                .build();
6214
6215            registry.register(tool);
6216
6217            // Both sessions should receive the notification
6218            let n1 = rx1.recv().await.unwrap();
6219            let n2 = rx2.recv().await.unwrap();
6220            assert!(matches!(n1, ServerNotification::ToolsListChanged));
6221            assert!(matches!(n2, ServerNotification::ToolsListChanged));
6222        }
6223
6224        #[tokio::test]
6225        async fn test_dynamic_tools_call_not_found() {
6226            let (router, _registry) = McpRouter::new()
6227                .server_info("test", "1.0")
6228                .with_dynamic_tools();
6229
6230            let mut router = router;
6231            init_router(&mut router).await;
6232
6233            let req = RouterRequest {
6234                id: RequestId::Number(1),
6235                inner: McpRequest::CallTool(CallToolParams {
6236                    name: "nonexistent".to_string(),
6237                    arguments: serde_json::json!({}),
6238                    meta: None,
6239                    task: None,
6240                }),
6241                extensions: Extensions::new(),
6242            };
6243
6244            let resp = router.ready().await.unwrap().call(req).await.unwrap();
6245            match resp.inner {
6246                Err(e) => {
6247                    assert_eq!(e.code, -32601);
6248                }
6249                _ => panic!("Expected method not found error"),
6250            }
6251        }
6252
6253        #[tokio::test]
6254        async fn test_dynamic_tools_registry_list() {
6255            let (_, registry) = McpRouter::new()
6256                .server_info("test", "1.0")
6257                .with_dynamic_tools();
6258
6259            assert!(registry.list().is_empty());
6260
6261            let tool = ToolBuilder::new("tool_a")
6262                .description("A")
6263                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6264                .build();
6265            registry.register(tool);
6266
6267            let tool = ToolBuilder::new("tool_b")
6268                .description("B")
6269                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6270                .build();
6271            registry.register(tool);
6272
6273            let tools = registry.list();
6274            assert_eq!(tools.len(), 2);
6275            let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
6276            assert!(names.contains(&"tool_a"));
6277            assert!(names.contains(&"tool_b"));
6278        }
6279    } // mod dynamic_tools_tests
6280
6281    #[tokio::test]
6282    async fn test_tool_if_true_registers() {
6283        let tool = ToolBuilder::new("conditional")
6284            .description("Conditional tool")
6285            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6286            .build();
6287
6288        let mut router = McpRouter::new().tool_if(true, tool);
6289        init_router(&mut router).await;
6290
6291        let req = RouterRequest {
6292            id: RequestId::Number(1),
6293            inner: McpRequest::ListTools(ListToolsParams::default()),
6294            extensions: Extensions::new(),
6295        };
6296        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6297        match resp.inner {
6298            Ok(McpResponse::ListTools(result)) => {
6299                assert_eq!(result.tools.len(), 1);
6300                assert_eq!(result.tools[0].name, "conditional");
6301            }
6302            _ => panic!("Expected ListTools response"),
6303        }
6304    }
6305
6306    #[tokio::test]
6307    async fn test_tool_if_false_skips() {
6308        let tool = ToolBuilder::new("conditional")
6309            .description("Conditional tool")
6310            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6311            .build();
6312
6313        let mut router = McpRouter::new().tool_if(false, tool);
6314        init_router(&mut router).await;
6315
6316        let req = RouterRequest {
6317            id: RequestId::Number(1),
6318            inner: McpRequest::ListTools(ListToolsParams::default()),
6319            extensions: Extensions::new(),
6320        };
6321        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6322        match resp.inner {
6323            Ok(McpResponse::ListTools(result)) => {
6324                assert_eq!(result.tools.len(), 0);
6325            }
6326            _ => panic!("Expected ListTools response"),
6327        }
6328    }
6329
6330    #[tokio::test]
6331    async fn test_tools_if_batch_conditional() {
6332        let tools = vec![
6333            ToolBuilder::new("a")
6334                .description("Tool A")
6335                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6336                .build(),
6337            ToolBuilder::new("b")
6338                .description("Tool B")
6339                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6340                .build(),
6341        ];
6342
6343        let mut router = McpRouter::new().tools_if(false, tools);
6344        init_router(&mut router).await;
6345
6346        let req = RouterRequest {
6347            id: RequestId::Number(1),
6348            inner: McpRequest::ListTools(ListToolsParams::default()),
6349            extensions: Extensions::new(),
6350        };
6351        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6352        match resp.inner {
6353            Ok(McpResponse::ListTools(result)) => {
6354                assert_eq!(result.tools.len(), 0);
6355            }
6356            _ => panic!("Expected ListTools response"),
6357        }
6358    }
6359
6360    #[test]
6361    fn test_resource_if_true_registers() {
6362        let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6363            .name("test")
6364            .text("hello");
6365
6366        let router = McpRouter::new().resource_if(true, resource);
6367        assert_eq!(router.inner.resources.len(), 1);
6368    }
6369
6370    #[test]
6371    fn test_resource_if_false_skips() {
6372        let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
6373            .name("test")
6374            .text("hello");
6375
6376        let router = McpRouter::new().resource_if(false, resource);
6377        assert_eq!(router.inner.resources.len(), 0);
6378    }
6379
6380    #[test]
6381    fn test_prompt_if_true_registers() {
6382        let prompt = crate::prompt::PromptBuilder::new("greet")
6383            .description("Greeting")
6384            .user_message("Hello!");
6385
6386        let router = McpRouter::new().prompt_if(true, prompt);
6387        assert_eq!(router.inner.prompts.len(), 1);
6388    }
6389
6390    #[test]
6391    fn test_prompt_if_false_skips() {
6392        let prompt = crate::prompt::PromptBuilder::new("greet")
6393            .description("Greeting")
6394            .user_message("Hello!");
6395
6396        let router = McpRouter::new().prompt_if(false, prompt);
6397        assert_eq!(router.inner.prompts.len(), 0);
6398    }
6399
6400    #[tokio::test]
6401    async fn test_disable_tool_hides_from_list() {
6402        let safe = ToolBuilder::new("safe")
6403            .description("Safe tool")
6404            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6405            .build();
6406        let dangerous = ToolBuilder::new("dangerous")
6407            .description("Dangerous tool")
6408            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6409            .build();
6410        let mut router = McpRouter::new().tool(safe).tool(dangerous);
6411        init_router(&mut router).await;
6412
6413        router.disable_tool("dangerous");
6414        assert!(router.is_tool_enabled("safe"));
6415        assert!(!router.is_tool_enabled("dangerous"));
6416
6417        let req = RouterRequest {
6418            id: RequestId::Number(1),
6419            inner: McpRequest::ListTools(ListToolsParams::default()),
6420            extensions: Extensions::new(),
6421        };
6422        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6423        match resp.inner {
6424            Ok(McpResponse::ListTools(result)) => {
6425                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
6426                assert_eq!(names, vec!["safe"]);
6427            }
6428            _ => panic!("Expected ListTools response"),
6429        }
6430    }
6431
6432    #[tokio::test]
6433    async fn test_disable_tool_blocks_call() {
6434        let dangerous = ToolBuilder::new("dangerous")
6435            .description("Dangerous tool")
6436            .handler(|_: AddInput| async { Ok(CallToolResult::text("ran")) })
6437            .build();
6438        let mut router = McpRouter::new().tool(dangerous);
6439        init_router(&mut router).await;
6440
6441        router.disable_tool("dangerous");
6442
6443        let req = RouterRequest {
6444            id: RequestId::Number(2),
6445            inner: McpRequest::CallTool(CallToolParams {
6446                name: "dangerous".to_string(),
6447                arguments: serde_json::json!({"a": 1, "b": 2}),
6448                meta: None,
6449                task: None,
6450            }),
6451            extensions: Extensions::new(),
6452        };
6453        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6454        let err = resp.inner.expect_err("disabled tool should error");
6455        assert_eq!(err.code, crate::error::ErrorCode::MethodNotFound as i32);
6456    }
6457
6458    #[tokio::test]
6459    async fn test_enable_tool_restores_visibility() {
6460        let tool = ToolBuilder::new("flippy")
6461            .description("Toggleable tool")
6462            .handler(|_: AddInput| async { Ok(CallToolResult::text("ran")) })
6463            .build();
6464        let mut router = McpRouter::new().tool(tool);
6465        init_router(&mut router).await;
6466
6467        router.disable_tool("flippy");
6468        router.enable_tool("flippy");
6469        assert!(router.is_tool_enabled("flippy"));
6470
6471        let req = RouterRequest {
6472            id: RequestId::Number(3),
6473            inner: McpRequest::CallTool(CallToolParams {
6474                name: "flippy".to_string(),
6475                arguments: serde_json::json!({"a": 1, "b": 2}),
6476                meta: None,
6477                task: None,
6478            }),
6479            extensions: Extensions::new(),
6480        };
6481        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6482        match resp.inner {
6483            Ok(McpResponse::CallTool(result)) => {
6484                assert_eq!(result.first_text(), Some("ran"));
6485            }
6486            _ => panic!("Expected CallTool response"),
6487        }
6488    }
6489
6490    #[tokio::test]
6491    async fn test_disable_propagates_through_fresh_session() {
6492        let tool = ToolBuilder::new("shared")
6493            .description("Shared across sessions")
6494            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
6495            .build();
6496        let router = McpRouter::new().tool(tool);
6497
6498        // Disable on the parent, observe via with_fresh_session clone.
6499        router.disable_tool("shared");
6500        let mut child = router.with_fresh_session();
6501        init_router(&mut child).await;
6502        assert!(!child.is_tool_enabled("shared"));
6503
6504        let req = RouterRequest {
6505            id: RequestId::Number(4),
6506            inner: McpRequest::ListTools(ListToolsParams::default()),
6507            extensions: Extensions::new(),
6508        };
6509        let resp = child.ready().await.unwrap().call(req).await.unwrap();
6510        match resp.inner {
6511            Ok(McpResponse::ListTools(result)) => {
6512                assert!(result.tools.is_empty());
6513            }
6514            _ => panic!("Expected ListTools response"),
6515        }
6516    }
6517
6518    #[tokio::test]
6519    async fn test_disable_resource_and_prompt() {
6520        let resource = crate::resource::ResourceBuilder::new("file:///hidden.txt")
6521            .name("hidden")
6522            .text("secret");
6523        let prompt = crate::prompt::PromptBuilder::new("hidden_prompt")
6524            .description("hidden")
6525            .user_message("hello");
6526
6527        let mut router = McpRouter::new().resource(resource).prompt(prompt);
6528        init_router(&mut router).await;
6529
6530        router.disable_resource("file:///hidden.txt");
6531        router.disable_prompt("hidden_prompt");
6532        assert!(!router.is_resource_enabled("file:///hidden.txt"));
6533        assert!(!router.is_prompt_enabled("hidden_prompt"));
6534
6535        // resources/list excludes
6536        let req = RouterRequest {
6537            id: RequestId::Number(5),
6538            inner: McpRequest::ListResources(ListResourcesParams::default()),
6539            extensions: Extensions::new(),
6540        };
6541        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6542        match resp.inner {
6543            Ok(McpResponse::ListResources(result)) => {
6544                assert!(result.resources.is_empty());
6545            }
6546            _ => panic!("Expected ListResources response"),
6547        }
6548
6549        // resources/read returns not found
6550        let req = RouterRequest {
6551            id: RequestId::Number(6),
6552            inner: McpRequest::ReadResource(ReadResourceParams {
6553                uri: "file:///hidden.txt".to_string(),
6554                meta: None,
6555            }),
6556            extensions: Extensions::new(),
6557        };
6558        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6559        let err = resp.inner.expect_err("disabled resource should error");
6560        assert_eq!(err.code, -32602); // SEP-2164: ResourceNotFound now uses InvalidParams
6561
6562        // prompts/list excludes
6563        let req = RouterRequest {
6564            id: RequestId::Number(7),
6565            inner: McpRequest::ListPrompts(ListPromptsParams::default()),
6566            extensions: Extensions::new(),
6567        };
6568        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6569        match resp.inner {
6570            Ok(McpResponse::ListPrompts(result)) => {
6571                assert!(result.prompts.is_empty());
6572            }
6573            _ => panic!("Expected ListPrompts response"),
6574        }
6575
6576        // prompts/get returns not found
6577        let req = RouterRequest {
6578            id: RequestId::Number(8),
6579            inner: McpRequest::GetPrompt(GetPromptParams {
6580                name: "hidden_prompt".to_string(),
6581                arguments: Default::default(),
6582                meta: None,
6583            }),
6584            extensions: Extensions::new(),
6585        };
6586        let resp = router.ready().await.unwrap().call(req).await.unwrap();
6587        let err = resp.inner.expect_err("disabled prompt should error");
6588        assert_eq!(err.code, crate::error::ErrorCode::MethodNotFound as i32);
6589    }
6590
6591    #[test]
6592    fn test_router_request_new() {
6593        let req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6594        assert_eq!(req.id, RequestId::Number(1));
6595        assert!(req.extensions.is_empty());
6596    }
6597
6598    #[test]
6599    fn test_with_inner_preserves_extensions() {
6600        let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6601        req.extensions.insert(42u32);
6602
6603        let rewritten = req.with_inner(McpRequest::ListTools(Default::default()));
6604        assert!(matches!(rewritten.inner, McpRequest::ListTools(_)));
6605        assert_eq!(rewritten.id, RequestId::Number(1));
6606        assert_eq!(rewritten.extensions.get::<u32>(), Some(&42));
6607    }
6608
6609    #[test]
6610    fn test_with_id_and_inner_preserves_extensions() {
6611        let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6612        req.extensions.insert(String::from("token-abc"));
6613
6614        let rewritten = req.with_id_and_inner(
6615            RequestId::Number(99),
6616            McpRequest::ListResources(Default::default()),
6617        );
6618        assert_eq!(rewritten.id, RequestId::Number(99));
6619        assert!(matches!(rewritten.inner, McpRequest::ListResources(_)));
6620        assert_eq!(
6621            rewritten.extensions.get::<String>(),
6622            Some(&String::from("token-abc"))
6623        );
6624    }
6625
6626    #[test]
6627    fn test_clone_with_inner_preserves_extensions() {
6628        let mut req = RouterRequest::new(RequestId::Number(1), McpRequest::Ping);
6629        req.extensions.insert(true);
6630
6631        let cloned = req.clone_with_inner(McpRequest::ListTools(Default::default()));
6632
6633        // Original still intact
6634        assert!(matches!(req.inner, McpRequest::Ping));
6635        assert_eq!(req.extensions.get::<bool>(), Some(&true));
6636
6637        // Clone has new inner but same extensions
6638        assert!(matches!(cloned.inner, McpRequest::ListTools(_)));
6639        assert_eq!(cloned.extensions.get::<bool>(), Some(&true));
6640    }
6641
6642    #[test]
6643    fn test_router_response_is_error() {
6644        let ok_resp = RouterResponse {
6645            id: RequestId::Number(1),
6646            inner: Ok(McpResponse::Pong(Default::default())),
6647        };
6648        assert!(!ok_resp.is_error());
6649
6650        let err_resp = RouterResponse {
6651            id: RequestId::Number(2),
6652            inner: Err(JsonRpcError::internal_error("boom")),
6653        };
6654        assert!(err_resp.is_error());
6655    }
6656
6657    #[test]
6658    fn test_extensions_len_and_is_empty() {
6659        let mut ext = Extensions::new();
6660        assert!(ext.is_empty());
6661        assert_eq!(ext.len(), 0);
6662
6663        ext.insert(42u32);
6664        assert!(!ext.is_empty());
6665        assert_eq!(ext.len(), 1);
6666
6667        ext.insert(String::from("hello"));
6668        assert_eq!(ext.len(), 2);
6669    }
6670
6671    #[test]
6672    fn test_router_response_serde_roundtrip() {
6673        // Success response
6674        let response = RouterResponse {
6675            id: RequestId::Number(1),
6676            inner: Ok(McpResponse::Empty(EmptyResult {})),
6677        };
6678        let json = serde_json::to_string(&response).unwrap();
6679        let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6680        assert_eq!(deserialized.id, RequestId::Number(1));
6681        assert!(!deserialized.is_error());
6682
6683        // Error response
6684        let response = RouterResponse {
6685            id: RequestId::String("req-2".into()),
6686            inner: Err(JsonRpcError::method_not_found("unknown")),
6687        };
6688        let json = serde_json::to_string(&response).unwrap();
6689        let deserialized: RouterResponse = serde_json::from_str(&json).unwrap();
6690        assert_eq!(deserialized.id, RequestId::String("req-2".into()));
6691        assert!(deserialized.is_error());
6692    }
6693
6694    // =========================================================================
6695    // Issue #872: McpRequest::Discover unit tests
6696    // Unit tests that exercise the router dispatch directly via JsonRpcService,
6697    // without going through the HTTP transport layer.
6698    // =========================================================================
6699
6700    #[tokio::test]
6701    async fn test_discover_dispatch_via_jsonrpc_service() {
6702        // server/discover must work without any prior initialize call.
6703        // The router does NOT require session initialization for this RPC.
6704        let router = McpRouter::new().server_info("unit-test-server", "4.2.0");
6705        let mut service = JsonRpcService::new(router);
6706
6707        let req = JsonRpcRequest::new(1, "server/discover");
6708        let resp = service.call_single(req).await.unwrap();
6709
6710        match resp {
6711            JsonRpcResponse::Result(r) => {
6712                // supportedVersions must be a non-empty array.
6713                let versions = r
6714                    .result
6715                    .get("supportedVersions")
6716                    .and_then(|v| v.as_array())
6717                    .expect("result.supportedVersions must be an array");
6718                assert!(!versions.is_empty(), "supportedVersions must not be empty");
6719
6720                // serverInfo.name must match what we configured.
6721                assert_eq!(
6722                    r.result["serverInfo"]["name"], "unit-test-server",
6723                    "serverInfo.name must match configured value"
6724                );
6725                assert_eq!(
6726                    r.result["serverInfo"]["version"], "4.2.0",
6727                    "serverInfo.version must match configured value"
6728                );
6729
6730                // server/discover must NOT include singular protocolVersion
6731                // (that field belongs to the initialize response shape).
6732                assert!(
6733                    r.result.get("protocolVersion").is_none(),
6734                    "server/discover must NOT include protocolVersion: {:?}",
6735                    r.result
6736                );
6737            }
6738            JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
6739            _ => panic!("unexpected response variant"),
6740        }
6741    }
6742
6743    #[tokio::test]
6744    async fn test_discover_does_not_require_initialization() {
6745        // server/discover works on a freshly created, un-initialized router.
6746        // No prior initialize call is made -- the session state is empty.
6747        let router = McpRouter::new().server_info("fresh-router", "1.0.0");
6748        let mut service = JsonRpcService::new(router);
6749
6750        let req = JsonRpcRequest::new(2, "server/discover");
6751        let resp = service.call_single(req).await.unwrap();
6752
6753        // Must succeed -- not return an error about missing session/initialization.
6754        assert!(
6755            !matches!(resp, JsonRpcResponse::Error(_)),
6756            "server/discover must not require initialization: {:?}",
6757            resp
6758        );
6759    }
6760}