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