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 = tool
1666                    .ok_or_else(|| Error::JsonRpc(JsonRpcError::method_not_found(&params.name)))?;
1667
1668                // Check tool filter if configured
1669                if let Some(filter) = &self.inner.tool_filter
1670                    && !filter.is_visible(&self.session, &tool)
1671                {
1672                    return Err(filter.denial_error(&params.name));
1673                }
1674
1675                if let Some(task_params) = params.task {
1676                    // Task-augmented request: validate task_support != Forbidden
1677                    if matches!(tool.task_support, TaskSupportMode::Forbidden) {
1678                        return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1679                            "Tool '{}' does not support async tasks",
1680                            params.name
1681                        ))));
1682                    }
1683
1684                    // Create the task
1685                    let (task_id, cancellation_token) = self.inner.task_store.create_task(
1686                        &params.name,
1687                        params.arguments.clone(),
1688                        task_params.ttl,
1689                    );
1690
1691                    tracing::info!(task_id = %task_id, tool = %params.name, "Created async task");
1692
1693                    // Create a context for the async task execution
1694                    let progress_token = params.meta.and_then(|m| m.progress_token);
1695                    let ctx = self.create_context(request_id, progress_token);
1696
1697                    // Spawn the task execution in the background
1698                    let task_store = self.inner.task_store.clone();
1699                    let tool = tool.clone();
1700                    let arguments = params.arguments;
1701                    let task_id_clone = task_id.clone();
1702
1703                    tokio::spawn(async move {
1704                        // Check for cancellation before starting
1705                        if cancellation_token.is_cancelled() {
1706                            tracing::debug!(task_id = %task_id_clone, "Task cancelled before execution");
1707                            return;
1708                        }
1709
1710                        // Execute the tool
1711                        let result = tool.call_with_context(ctx, arguments).await;
1712
1713                        if cancellation_token.is_cancelled() {
1714                            tracing::debug!(task_id = %task_id_clone, "Task cancelled during execution");
1715                        } else if result.is_error {
1716                            // Tool returned an error result
1717                            let error_msg = result.first_text().unwrap_or("Tool execution failed");
1718                            task_store.fail_task(&task_id_clone, error_msg);
1719                            tracing::warn!(task_id = %task_id_clone, error = %error_msg, "Task failed");
1720                        } else {
1721                            task_store.complete_task(&task_id_clone, result);
1722                            tracing::debug!(task_id = %task_id_clone, "Task completed successfully");
1723                        }
1724                    });
1725
1726                    let task = self.inner.task_store.get_task(&task_id).ok_or_else(|| {
1727                        Error::JsonRpc(JsonRpcError::internal_error(
1728                            "Failed to retrieve created task",
1729                        ))
1730                    })?;
1731
1732                    Ok(McpResponse::CreateTask(CreateTaskResult {
1733                        task,
1734                        meta: None,
1735                    }))
1736                } else {
1737                    // Synchronous request: validate task_support != Required
1738                    if matches!(tool.task_support, TaskSupportMode::Required) {
1739                        return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
1740                            "Tool '{}' requires async task execution (include 'task' in params)",
1741                            params.name
1742                        ))));
1743                    }
1744
1745                    // Extract progress token from request metadata
1746                    let progress_token = params.meta.and_then(|m| m.progress_token);
1747                    let ctx = self.create_context(request_id, progress_token);
1748
1749                    tracing::debug!(tool = %params.name, "Calling tool");
1750                    let result = tool.call_with_context(ctx, params.arguments).await;
1751
1752                    Ok(McpResponse::CallTool(result))
1753                }
1754            }
1755
1756            McpRequest::ListResources(params) => {
1757                let is_visible = |r: &Resource| -> bool {
1758                    self.inner
1759                        .resource_filter
1760                        .as_ref()
1761                        .map(|f| f.is_visible(&self.session, r))
1762                        .unwrap_or(true)
1763                };
1764
1765                let mut resources: Vec<ResourceDefinition> = self
1766                    .inner
1767                    .resources
1768                    .values()
1769                    .filter(|r| is_visible(r))
1770                    .map(|r| r.definition())
1771                    .collect();
1772
1773                // Merge dynamic resources (static resources win on URI collision)
1774                #[cfg(feature = "dynamic-tools")]
1775                if let Some(ref dynamic) = self.inner.dynamic_resources {
1776                    let static_uris: HashSet<String> =
1777                        resources.iter().map(|r| r.uri.clone()).collect();
1778                    for r in dynamic.list() {
1779                        if !static_uris.contains(&r.uri) && is_visible(&r) {
1780                            resources.push(r.definition());
1781                        }
1782                    }
1783                }
1784
1785                resources.sort_by(|a, b| a.uri.cmp(&b.uri));
1786
1787                let (resources, next_cursor) =
1788                    paginate(resources, params.cursor.as_deref(), self.inner.page_size)?;
1789
1790                Ok(McpResponse::ListResources(ListResourcesResult {
1791                    resources,
1792                    next_cursor,
1793                    meta: None,
1794                }))
1795            }
1796
1797            McpRequest::ListResourceTemplates(params) => {
1798                let mut resource_templates: Vec<ResourceTemplateDefinition> = self
1799                    .inner
1800                    .resource_templates
1801                    .iter()
1802                    .map(|t| t.definition())
1803                    .collect();
1804
1805                // Merge dynamic resource templates (static win on collision)
1806                #[cfg(feature = "dynamic-tools")]
1807                if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
1808                    let static_patterns: HashSet<String> = resource_templates
1809                        .iter()
1810                        .map(|t| t.uri_template.clone())
1811                        .collect();
1812                    for t in dynamic.list() {
1813                        if !static_patterns.contains(&t.uri_template) {
1814                            resource_templates.push(t.definition());
1815                        }
1816                    }
1817                }
1818
1819                resource_templates.sort_by(|a, b| a.uri_template.cmp(&b.uri_template));
1820
1821                let (resource_templates, next_cursor) = paginate(
1822                    resource_templates,
1823                    params.cursor.as_deref(),
1824                    self.inner.page_size,
1825                )?;
1826
1827                Ok(McpResponse::ListResourceTemplates(
1828                    ListResourceTemplatesResult {
1829                        resource_templates,
1830                        next_cursor,
1831                        meta: None,
1832                    },
1833                ))
1834            }
1835
1836            McpRequest::ReadResource(params) => {
1837                // First, try to find a static resource
1838                if let Some(resource) = self.inner.resources.get(&params.uri) {
1839                    // Check resource filter if configured
1840                    if let Some(filter) = &self.inner.resource_filter
1841                        && !filter.is_visible(&self.session, resource)
1842                    {
1843                        return Err(filter.denial_error(&params.uri));
1844                    }
1845
1846                    tracing::debug!(uri = %params.uri, "Reading static resource");
1847                    let result = resource.read().await;
1848                    return Ok(McpResponse::ReadResource(result));
1849                }
1850
1851                // Try dynamic resources
1852                #[cfg(feature = "dynamic-tools")]
1853                #[allow(clippy::collapsible_if)]
1854                if let Some(ref dynamic) = self.inner.dynamic_resources {
1855                    if let Some(resource) = dynamic.get(&params.uri) {
1856                        if let Some(filter) = &self.inner.resource_filter
1857                            && !filter.is_visible(&self.session, &resource)
1858                        {
1859                            return Err(filter.denial_error(&params.uri));
1860                        }
1861                        tracing::debug!(uri = %params.uri, "Reading dynamic resource");
1862                        let result = resource.read().await;
1863                        return Ok(McpResponse::ReadResource(result));
1864                    }
1865                }
1866
1867                // Try static templates
1868                for template in &self.inner.resource_templates {
1869                    if let Some(variables) = template.match_uri(&params.uri) {
1870                        tracing::debug!(
1871                            uri = %params.uri,
1872                            template = %template.uri_template,
1873                            "Reading resource via template"
1874                        );
1875                        let result = template.read(&params.uri, variables).await?;
1876                        return Ok(McpResponse::ReadResource(result));
1877                    }
1878                }
1879
1880                // Try dynamic templates
1881                #[cfg(feature = "dynamic-tools")]
1882                #[allow(clippy::collapsible_if)]
1883                if let Some(ref dynamic) = self.inner.dynamic_resource_templates {
1884                    if let Some((template, variables)) = dynamic.match_uri(&params.uri) {
1885                        tracing::debug!(
1886                            uri = %params.uri,
1887                            template = %template.uri_template,
1888                            "Reading resource via dynamic template"
1889                        );
1890                        let result = template.read(&params.uri, variables).await?;
1891                        return Ok(McpResponse::ReadResource(result));
1892                    }
1893                }
1894
1895                // No match found
1896                Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1897                    &params.uri,
1898                )))
1899            }
1900
1901            McpRequest::SubscribeResource(params) => {
1902                // Verify the resource exists
1903                if !self.inner.resources.contains_key(&params.uri) {
1904                    return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1905                        &params.uri,
1906                    )));
1907                }
1908
1909                tracing::debug!(uri = %params.uri, "Subscribing to resource");
1910                self.subscribe(&params.uri);
1911
1912                Ok(McpResponse::SubscribeResource(EmptyResult {}))
1913            }
1914
1915            McpRequest::UnsubscribeResource(params) => {
1916                // Verify the resource exists
1917                if !self.inner.resources.contains_key(&params.uri) {
1918                    return Err(Error::JsonRpc(JsonRpcError::resource_not_found(
1919                        &params.uri,
1920                    )));
1921                }
1922
1923                tracing::debug!(uri = %params.uri, "Unsubscribing from resource");
1924                self.unsubscribe(&params.uri);
1925
1926                Ok(McpResponse::UnsubscribeResource(EmptyResult {}))
1927            }
1928
1929            McpRequest::ListPrompts(params) => {
1930                let is_visible = |p: &Prompt| -> bool {
1931                    self.inner
1932                        .prompt_filter
1933                        .as_ref()
1934                        .map(|f| f.is_visible(&self.session, p))
1935                        .unwrap_or(true)
1936                };
1937
1938                let mut prompts: Vec<PromptDefinition> = self
1939                    .inner
1940                    .prompts
1941                    .values()
1942                    .filter(|p| is_visible(p))
1943                    .map(|p| p.definition())
1944                    .collect();
1945
1946                // Merge dynamic prompts (static prompts win on name collision)
1947                #[cfg(feature = "dynamic-tools")]
1948                if let Some(ref dynamic) = self.inner.dynamic_prompts {
1949                    let static_names: HashSet<String> =
1950                        prompts.iter().map(|p| p.name.clone()).collect();
1951                    for p in dynamic.list() {
1952                        if !static_names.contains(&p.name) && is_visible(&p) {
1953                            prompts.push(p.definition());
1954                        }
1955                    }
1956                }
1957
1958                prompts.sort_by(|a, b| a.name.cmp(&b.name));
1959
1960                let (prompts, next_cursor) =
1961                    paginate(prompts, params.cursor.as_deref(), self.inner.page_size)?;
1962
1963                Ok(McpResponse::ListPrompts(ListPromptsResult {
1964                    prompts,
1965                    next_cursor,
1966                    meta: None,
1967                }))
1968            }
1969
1970            McpRequest::GetPrompt(params) => {
1971                // Look up static prompts first, then dynamic
1972                let prompt = self.inner.prompts.get(&params.name).cloned();
1973                #[cfg(feature = "dynamic-tools")]
1974                let prompt = prompt.or_else(|| {
1975                    self.inner
1976                        .dynamic_prompts
1977                        .as_ref()
1978                        .and_then(|d| d.get(&params.name))
1979                });
1980                let prompt = prompt.ok_or_else(|| {
1981                    Error::JsonRpc(JsonRpcError::method_not_found(&format!(
1982                        "Prompt not found: {}",
1983                        params.name
1984                    )))
1985                })?;
1986
1987                // Check prompt filter if configured
1988                if let Some(filter) = &self.inner.prompt_filter
1989                    && !filter.is_visible(&self.session, &prompt)
1990                {
1991                    return Err(filter.denial_error(&params.name));
1992                }
1993
1994                tracing::debug!(name = %params.name, "Getting prompt");
1995                let result = prompt.get(params.arguments).await?;
1996
1997                Ok(McpResponse::GetPrompt(result))
1998            }
1999
2000            McpRequest::Ping => Ok(McpResponse::Pong(EmptyResult {})),
2001
2002            McpRequest::ListTasks(params) => {
2003                let tasks = self.inner.task_store.list_tasks(params.status);
2004
2005                let (tasks, next_cursor) =
2006                    paginate(tasks, params.cursor.as_deref(), self.inner.page_size)?;
2007
2008                Ok(McpResponse::ListTasks(ListTasksResult {
2009                    tasks,
2010                    next_cursor,
2011                }))
2012            }
2013
2014            McpRequest::GetTaskInfo(params) => {
2015                let task = self
2016                    .inner
2017                    .task_store
2018                    .get_task(&params.task_id)
2019                    .ok_or_else(|| {
2020                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2021                            "Task not found: {}",
2022                            params.task_id
2023                        )))
2024                    })?;
2025
2026                Ok(McpResponse::GetTaskInfo(task))
2027            }
2028
2029            McpRequest::GetTaskResult(params) => {
2030                // Wait for task to reach terminal state (blocks if still running)
2031                let (task_obj, result, error) = self
2032                    .inner
2033                    .task_store
2034                    .wait_for_completion(&params.task_id)
2035                    .await
2036                    .ok_or_else(|| {
2037                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2038                            "Task not found: {}",
2039                            params.task_id
2040                        )))
2041                    })?;
2042
2043                // Build _meta with related-task reference
2044                let meta = serde_json::json!({
2045                    "io.modelcontextprotocol/related-task": task_obj
2046                });
2047
2048                match task_obj.status {
2049                    TaskStatus::Cancelled => Err(Error::JsonRpc(JsonRpcError::invalid_params(
2050                        format!("Task {} was cancelled", params.task_id),
2051                    ))),
2052                    TaskStatus::Failed => {
2053                        let mut call_result = CallToolResult::error(
2054                            error.unwrap_or_else(|| "Task failed".to_string()),
2055                        );
2056                        call_result.meta = Some(meta);
2057                        Ok(McpResponse::GetTaskResult(call_result))
2058                    }
2059                    _ => {
2060                        let mut call_result = result.unwrap_or_else(|| CallToolResult::text(""));
2061                        call_result.meta = Some(meta);
2062                        Ok(McpResponse::GetTaskResult(call_result))
2063                    }
2064                }
2065            }
2066
2067            McpRequest::CancelTask(params) => {
2068                // First check if the task exists and is not already terminal
2069                let current = self
2070                    .inner
2071                    .task_store
2072                    .get_task(&params.task_id)
2073                    .ok_or_else(|| {
2074                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2075                            "Task not found: {}",
2076                            params.task_id
2077                        )))
2078                    })?;
2079
2080                if current.status.is_terminal() {
2081                    return Err(Error::JsonRpc(JsonRpcError::invalid_params(format!(
2082                        "Task {} is already in terminal state: {}",
2083                        params.task_id, current.status
2084                    ))));
2085                }
2086
2087                let task_obj = self
2088                    .inner
2089                    .task_store
2090                    .cancel_task(&params.task_id, params.reason.as_deref())
2091                    .ok_or_else(|| {
2092                        Error::JsonRpc(JsonRpcError::invalid_params(format!(
2093                            "Task not found: {}",
2094                            params.task_id
2095                        )))
2096                    })?;
2097
2098                Ok(McpResponse::CancelTask(task_obj))
2099            }
2100
2101            McpRequest::SetLoggingLevel(params) => {
2102                tracing::debug!(level = ?params.level, "Client set logging level");
2103                if let Ok(mut level) = self.inner.min_log_level.write() {
2104                    *level = params.level;
2105                }
2106                Ok(McpResponse::SetLoggingLevel(EmptyResult {}))
2107            }
2108
2109            McpRequest::Complete(params) => {
2110                tracing::debug!(
2111                    reference = ?params.reference,
2112                    argument = %params.argument.name,
2113                    "Completion request"
2114                );
2115
2116                // Delegate to registered completion handler if available
2117                if let Some(ref handler) = self.inner.completion_handler {
2118                    let result = handler(params).await?;
2119                    Ok(McpResponse::Complete(result))
2120                } else {
2121                    // No completion handler registered, return empty completions
2122                    Ok(McpResponse::Complete(CompleteResult::new(vec![])))
2123                }
2124            }
2125
2126            McpRequest::Unknown { method, .. } => {
2127                Err(Error::JsonRpc(JsonRpcError::method_not_found(&method)))
2128            }
2129            _ => Err(Error::JsonRpc(JsonRpcError::method_not_found(
2130                "unknown method",
2131            ))),
2132        }
2133    }
2134
2135    /// Handle an MCP notification (no response expected)
2136    pub fn handle_notification(&self, notification: McpNotification) {
2137        match notification {
2138            McpNotification::Initialized => {
2139                let phase_before = self.session.phase();
2140                if self.session.mark_initialized() {
2141                    if phase_before == crate::session::SessionPhase::Uninitialized {
2142                        tracing::info!(
2143                            "Session initialized from uninitialized state (race resolved)"
2144                        );
2145                    } else {
2146                        tracing::info!("Session initialized, entering operation phase");
2147                    }
2148                } else {
2149                    tracing::warn!(
2150                        phase = ?self.session.phase(),
2151                        "Received initialized notification in unexpected state"
2152                    );
2153                }
2154            }
2155            McpNotification::Cancelled(params) => {
2156                if let Some(ref request_id) = params.request_id {
2157                    if self.cancel_request(request_id) {
2158                        tracing::info!(
2159                            request_id = ?request_id,
2160                            reason = ?params.reason,
2161                            "Request cancelled"
2162                        );
2163                    } else {
2164                        tracing::debug!(
2165                            request_id = ?request_id,
2166                            reason = ?params.reason,
2167                            "Cancellation requested for unknown request"
2168                        );
2169                    }
2170                } else {
2171                    tracing::debug!(
2172                        reason = ?params.reason,
2173                        "Cancellation notification received without request_id"
2174                    );
2175                }
2176            }
2177            McpNotification::Progress(params) => {
2178                tracing::trace!(
2179                    token = ?params.progress_token,
2180                    progress = params.progress,
2181                    total = ?params.total,
2182                    "Progress notification"
2183                );
2184                // Progress notifications from client are unusual but valid
2185            }
2186            McpNotification::RootsListChanged => {
2187                tracing::info!("Client roots list changed");
2188                // Server should re-request roots if needed
2189                // This is handled by the application layer
2190            }
2191            McpNotification::Unknown { method, .. } => {
2192                tracing::debug!(method = %method, "Unknown notification received");
2193            }
2194            _ => {
2195                tracing::debug!("Unrecognized notification variant received");
2196            }
2197        }
2198    }
2199}
2200
2201impl Default for McpRouter {
2202    fn default() -> Self {
2203        Self::new()
2204    }
2205}
2206
2207// =============================================================================
2208// Tower Service implementation
2209// =============================================================================
2210
2211// Re-export Extensions from context for backwards compatibility
2212pub use crate::context::Extensions;
2213
2214/// A map of tool names to their annotations, for use by middleware.
2215///
2216/// This is automatically inserted into [`RouterRequest::extensions`] for
2217/// `tools/call` requests, allowing middleware to inspect tool safety hints
2218/// (e.g., `read_only_hint`, `destructive_hint`) without needing direct
2219/// access to the router's tool registry.
2220///
2221/// # Example
2222///
2223/// ```rust,ignore
2224/// use tower_mcp::router::ToolAnnotationsMap;
2225/// use tower_mcp::protocol::McpRequest;
2226///
2227/// // In a middleware Service::call():
2228/// fn call(&mut self, req: RouterRequest) -> Self::Future {
2229///     if let McpRequest::CallTool(params) = &req.inner {
2230///         if let Some(map) = req.extensions.get::<ToolAnnotationsMap>() {
2231///             let annotations = map.get(&params.name);
2232///             // Check annotations.read_only_hint, destructive_hint, etc.
2233///         }
2234///     }
2235///     self.inner.call(req)
2236/// }
2237/// ```
2238#[derive(Debug, Clone)]
2239pub struct ToolAnnotationsMap {
2240    map: Arc<HashMap<String, ToolAnnotations>>,
2241}
2242
2243impl ToolAnnotationsMap {
2244    /// Look up annotations for a tool by name.
2245    ///
2246    /// Returns `None` if the tool has no annotations or doesn't exist.
2247    pub fn get(&self, tool_name: &str) -> Option<&ToolAnnotations> {
2248        self.map.get(tool_name)
2249    }
2250
2251    /// Check if a tool is read-only (does not modify state).
2252    ///
2253    /// Returns `false` if the tool has no annotations or doesn't exist
2254    /// (the MCP spec default for `readOnlyHint` is `false`).
2255    pub fn is_read_only(&self, tool_name: &str) -> bool {
2256        self.map.get(tool_name).is_some_and(|a| a.read_only_hint)
2257    }
2258
2259    /// Check if a tool may have destructive effects.
2260    ///
2261    /// Returns `true` if the tool has no annotations or doesn't exist
2262    /// (the MCP spec default for `destructiveHint` is `true`).
2263    pub fn is_destructive(&self, tool_name: &str) -> bool {
2264        self.map.get(tool_name).is_none_or(|a| a.destructive_hint)
2265    }
2266
2267    /// Check if a tool is idempotent.
2268    ///
2269    /// Returns `false` if the tool has no annotations or doesn't exist
2270    /// (the MCP spec default for `idempotentHint` is `false`).
2271    pub fn is_idempotent(&self, tool_name: &str) -> bool {
2272        self.map.get(tool_name).is_some_and(|a| a.idempotent_hint)
2273    }
2274}
2275
2276/// Request type for the tower Service implementation
2277#[derive(Debug, Clone)]
2278pub struct RouterRequest {
2279    /// The JSON-RPC request ID.
2280    pub id: RequestId,
2281    /// The parsed MCP request.
2282    pub inner: McpRequest,
2283    /// Type-map for passing data (e.g., `TokenClaims`) through middleware.
2284    pub extensions: Extensions,
2285}
2286
2287/// Response type for the tower Service implementation
2288#[derive(Debug, Clone)]
2289pub struct RouterResponse {
2290    /// The JSON-RPC request ID this response corresponds to.
2291    pub id: RequestId,
2292    /// The MCP response or JSON-RPC error.
2293    pub inner: std::result::Result<McpResponse, JsonRpcError>,
2294}
2295
2296impl RouterResponse {
2297    /// Convert to JSON-RPC response
2298    pub fn into_jsonrpc(self) -> JsonRpcResponse {
2299        match self.inner {
2300            Ok(response) => match serde_json::to_value(response) {
2301                Ok(result) => JsonRpcResponse::result(self.id, result),
2302                Err(e) => {
2303                    tracing::error!(error = %e, "Failed to serialize response");
2304                    JsonRpcResponse::error(
2305                        Some(self.id),
2306                        JsonRpcError::internal_error(format!("Serialization error: {}", e)),
2307                    )
2308                }
2309            },
2310            Err(error) => JsonRpcResponse::error(Some(self.id), error),
2311        }
2312    }
2313}
2314
2315impl Service<RouterRequest> for McpRouter {
2316    type Response = RouterResponse;
2317    type Error = std::convert::Infallible; // Errors are in the response
2318    type Future =
2319        Pin<Box<dyn Future<Output = std::result::Result<Self::Response, Self::Error>> + Send>>;
2320
2321    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
2322        Poll::Ready(Ok(()))
2323    }
2324
2325    fn call(&mut self, req: RouterRequest) -> Self::Future {
2326        let router = self.clone();
2327        let request_id = req.id.clone();
2328        Box::pin(async move {
2329            let result = router.handle(req.id, req.inner).await;
2330            // Clean up tracking after request completes
2331            router.complete_request(&request_id);
2332            Ok(RouterResponse {
2333                id: request_id,
2334                // Map tower-mcp errors to JSON-RPC errors:
2335                // - Error::JsonRpc: forwarded as-is (preserves original code)
2336                // - Error::Tool: mapped to -32603 (Internal Error)
2337                // - All others: mapped to -32603 (Internal Error)
2338                inner: result.map_err(|e| match e {
2339                    Error::JsonRpc(err) => err,
2340                    Error::Tool(err) => JsonRpcError::internal_error(err.to_string()),
2341                    e => JsonRpcError::internal_error(e.to_string()),
2342                }),
2343            })
2344        })
2345    }
2346}
2347
2348#[cfg(test)]
2349mod tests {
2350    use super::*;
2351    use crate::extract::{Context, Json};
2352    use crate::jsonrpc::JsonRpcService;
2353    use crate::tool::ToolBuilder;
2354    use schemars::JsonSchema;
2355    use serde::Deserialize;
2356    use tower::ServiceExt;
2357
2358    #[derive(Debug, Deserialize, JsonSchema)]
2359    struct AddInput {
2360        a: i64,
2361        b: i64,
2362    }
2363
2364    /// Helper to initialize a router for testing
2365    async fn init_router(router: &mut McpRouter) {
2366        // Send initialize request
2367        let init_req = RouterRequest {
2368            id: RequestId::Number(0),
2369            inner: McpRequest::Initialize(InitializeParams {
2370                protocol_version: "2025-11-25".to_string(),
2371                capabilities: ClientCapabilities {
2372                    roots: None,
2373                    sampling: None,
2374                    elicitation: None,
2375                    tasks: None,
2376                    experimental: None,
2377                    extensions: None,
2378                },
2379                client_info: Implementation {
2380                    name: "test".to_string(),
2381                    version: "1.0".to_string(),
2382                    ..Default::default()
2383                },
2384                meta: None,
2385            }),
2386            extensions: Extensions::new(),
2387        };
2388        let _ = router.ready().await.unwrap().call(init_req).await.unwrap();
2389        // Send initialized notification
2390        router.handle_notification(McpNotification::Initialized);
2391    }
2392
2393    #[tokio::test]
2394    async fn test_router_list_tools() {
2395        let add_tool = ToolBuilder::new("add")
2396            .description("Add two numbers")
2397            .handler(|input: AddInput| async move {
2398                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2399            })
2400            .build();
2401
2402        let mut router = McpRouter::new().tool(add_tool);
2403
2404        // Initialize session first
2405        init_router(&mut router).await;
2406
2407        let req = RouterRequest {
2408            id: RequestId::Number(1),
2409            inner: McpRequest::ListTools(ListToolsParams::default()),
2410            extensions: Extensions::new(),
2411        };
2412
2413        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2414
2415        match resp.inner {
2416            Ok(McpResponse::ListTools(result)) => {
2417                assert_eq!(result.tools.len(), 1);
2418                assert_eq!(result.tools[0].name, "add");
2419            }
2420            _ => panic!("Expected ListTools response"),
2421        }
2422    }
2423
2424    #[tokio::test]
2425    async fn test_router_call_tool() {
2426        let add_tool = ToolBuilder::new("add")
2427            .description("Add two numbers")
2428            .handler(|input: AddInput| async move {
2429                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2430            })
2431            .build();
2432
2433        let mut router = McpRouter::new().tool(add_tool);
2434
2435        // Initialize session first
2436        init_router(&mut router).await;
2437
2438        let req = RouterRequest {
2439            id: RequestId::Number(1),
2440            inner: McpRequest::CallTool(CallToolParams {
2441                name: "add".to_string(),
2442                arguments: serde_json::json!({"a": 2, "b": 3}),
2443                meta: None,
2444                task: None,
2445            }),
2446            extensions: Extensions::new(),
2447        };
2448
2449        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2450
2451        match resp.inner {
2452            Ok(McpResponse::CallTool(result)) => {
2453                assert!(!result.is_error);
2454                // Check the text content
2455                match &result.content[0] {
2456                    Content::Text { text, .. } => assert_eq!(text, "5"),
2457                    _ => panic!("Expected text content"),
2458                }
2459            }
2460            _ => panic!("Expected CallTool response"),
2461        }
2462    }
2463
2464    /// Helper to initialize a JsonRpcService for testing
2465    async fn init_jsonrpc_service(service: &mut JsonRpcService<McpRouter>, router: &McpRouter) {
2466        let init_req = JsonRpcRequest::new(0, "initialize").with_params(serde_json::json!({
2467            "protocolVersion": "2025-11-25",
2468            "capabilities": {},
2469            "clientInfo": { "name": "test", "version": "1.0" }
2470        }));
2471        let _ = service.call_single(init_req).await.unwrap();
2472        router.handle_notification(McpNotification::Initialized);
2473    }
2474
2475    #[tokio::test]
2476    async fn test_jsonrpc_service() {
2477        let add_tool = ToolBuilder::new("add")
2478            .description("Add two numbers")
2479            .handler(|input: AddInput| async move {
2480                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2481            })
2482            .build();
2483
2484        let router = McpRouter::new().tool(add_tool);
2485        let mut service = JsonRpcService::new(router.clone());
2486
2487        // Initialize session first
2488        init_jsonrpc_service(&mut service, &router).await;
2489
2490        let req = JsonRpcRequest::new(1, "tools/list");
2491
2492        let resp = service.call_single(req).await.unwrap();
2493
2494        match resp {
2495            JsonRpcResponse::Result(r) => {
2496                assert_eq!(r.id, RequestId::Number(1));
2497                let tools = r.result.get("tools").unwrap().as_array().unwrap();
2498                assert_eq!(tools.len(), 1);
2499            }
2500            JsonRpcResponse::Error(_) => panic!("Expected success response"),
2501            _ => panic!("unexpected response variant"),
2502        }
2503    }
2504
2505    #[tokio::test]
2506    async fn test_batch_request() {
2507        let add_tool = ToolBuilder::new("add")
2508            .description("Add two numbers")
2509            .handler(|input: AddInput| async move {
2510                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2511            })
2512            .build();
2513
2514        let router = McpRouter::new().tool(add_tool);
2515        let mut service = JsonRpcService::new(router.clone());
2516
2517        // Initialize session first
2518        init_jsonrpc_service(&mut service, &router).await;
2519
2520        // Create a batch of requests
2521        let requests = vec![
2522            JsonRpcRequest::new(1, "tools/list"),
2523            JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2524                "name": "add",
2525                "arguments": {"a": 10, "b": 20}
2526            })),
2527            JsonRpcRequest::new(3, "ping"),
2528        ];
2529
2530        let responses = service.call_batch(requests).await.unwrap();
2531
2532        assert_eq!(responses.len(), 3);
2533
2534        // Check first response (tools/list)
2535        match &responses[0] {
2536            JsonRpcResponse::Result(r) => {
2537                assert_eq!(r.id, RequestId::Number(1));
2538                let tools = r.result.get("tools").unwrap().as_array().unwrap();
2539                assert_eq!(tools.len(), 1);
2540            }
2541            JsonRpcResponse::Error(_) => panic!("Expected success for tools/list"),
2542            _ => panic!("unexpected response variant"),
2543        }
2544
2545        // Check second response (tools/call)
2546        match &responses[1] {
2547            JsonRpcResponse::Result(r) => {
2548                assert_eq!(r.id, RequestId::Number(2));
2549                let content = r.result.get("content").unwrap().as_array().unwrap();
2550                let text = content[0].get("text").unwrap().as_str().unwrap();
2551                assert_eq!(text, "30");
2552            }
2553            JsonRpcResponse::Error(_) => panic!("Expected success for tools/call"),
2554            _ => panic!("unexpected response variant"),
2555        }
2556
2557        // Check third response (ping)
2558        match &responses[2] {
2559            JsonRpcResponse::Result(r) => {
2560                assert_eq!(r.id, RequestId::Number(3));
2561            }
2562            JsonRpcResponse::Error(_) => panic!("Expected success for ping"),
2563            _ => panic!("unexpected response variant"),
2564        }
2565    }
2566
2567    #[tokio::test]
2568    async fn test_empty_batch_error() {
2569        let router = McpRouter::new();
2570        let mut service = JsonRpcService::new(router);
2571
2572        let result = service.call_batch(vec![]).await;
2573        assert!(result.is_err());
2574    }
2575
2576    // =========================================================================
2577    // Progress Token Tests
2578    // =========================================================================
2579
2580    #[tokio::test]
2581    async fn test_progress_token_extraction() {
2582        use crate::context::{ServerNotification, notification_channel};
2583        use crate::protocol::ProgressToken;
2584        use std::sync::Arc;
2585        use std::sync::atomic::{AtomicBool, Ordering};
2586
2587        // Track whether progress was reported
2588        let progress_reported = Arc::new(AtomicBool::new(false));
2589        let progress_ref = progress_reported.clone();
2590
2591        // Create a tool that reports progress
2592        let tool = ToolBuilder::new("progress_tool")
2593            .description("Tool that reports progress")
2594            .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2595                let reported = progress_ref.clone();
2596                async move {
2597                    // Report progress - this should work if token was extracted
2598                    ctx.report_progress(50.0, Some(100.0), Some("Halfway"))
2599                        .await;
2600                    reported.store(true, Ordering::SeqCst);
2601                    Ok(CallToolResult::text("done"))
2602                }
2603            })
2604            .build();
2605
2606        // Set up notification channel
2607        let (tx, mut rx) = notification_channel(10);
2608        let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2609        let mut service = JsonRpcService::new(router.clone());
2610
2611        // Initialize
2612        init_jsonrpc_service(&mut service, &router).await;
2613
2614        // Call tool WITH progress token in _meta
2615        let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2616            "name": "progress_tool",
2617            "arguments": {"a": 1, "b": 2},
2618            "_meta": {
2619                "progressToken": "test-token-123"
2620            }
2621        }));
2622
2623        let resp = service.call_single(req).await.unwrap();
2624
2625        // Verify the tool was called successfully
2626        match resp {
2627            JsonRpcResponse::Result(_) => {}
2628            JsonRpcResponse::Error(e) => panic!("Expected success, got error: {:?}", e),
2629            _ => panic!("unexpected response variant"),
2630        }
2631
2632        // Verify progress was reported by handler
2633        assert!(progress_reported.load(Ordering::SeqCst));
2634
2635        // Verify progress notification was sent through channel
2636        let notification = rx.try_recv().expect("Expected progress notification");
2637        match notification {
2638            ServerNotification::Progress(params) => {
2639                assert_eq!(
2640                    params.progress_token,
2641                    ProgressToken::String("test-token-123".to_string())
2642                );
2643                assert_eq!(params.progress, 50.0);
2644                assert_eq!(params.total, Some(100.0));
2645                assert_eq!(params.message.as_deref(), Some("Halfway"));
2646            }
2647            _ => panic!("Expected Progress notification"),
2648        }
2649    }
2650
2651    #[tokio::test]
2652    async fn test_tool_call_without_progress_token() {
2653        use crate::context::notification_channel;
2654        use std::sync::Arc;
2655        use std::sync::atomic::{AtomicBool, Ordering};
2656
2657        let progress_attempted = Arc::new(AtomicBool::new(false));
2658        let progress_ref = progress_attempted.clone();
2659
2660        let tool = ToolBuilder::new("no_token_tool")
2661            .description("Tool that tries to report progress without token")
2662            .extractor_handler((), move |ctx: Context, Json(_input): Json<AddInput>| {
2663                let attempted = progress_ref.clone();
2664                async move {
2665                    // Try to report progress - should be a no-op without token
2666                    ctx.report_progress(50.0, Some(100.0), None).await;
2667                    attempted.store(true, Ordering::SeqCst);
2668                    Ok(CallToolResult::text("done"))
2669                }
2670            })
2671            .build();
2672
2673        let (tx, mut rx) = notification_channel(10);
2674        let router = McpRouter::new().with_notification_sender(tx).tool(tool);
2675        let mut service = JsonRpcService::new(router.clone());
2676
2677        init_jsonrpc_service(&mut service, &router).await;
2678
2679        // Call tool WITHOUT progress token
2680        let req = JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2681            "name": "no_token_tool",
2682            "arguments": {"a": 1, "b": 2}
2683        }));
2684
2685        let resp = service.call_single(req).await.unwrap();
2686        assert!(matches!(resp, JsonRpcResponse::Result(_)));
2687
2688        // Handler was called
2689        assert!(progress_attempted.load(Ordering::SeqCst));
2690
2691        // But no notification was sent (no progress token)
2692        assert!(rx.try_recv().is_err());
2693    }
2694
2695    #[tokio::test]
2696    async fn test_batch_errors_returned_not_dropped() {
2697        let add_tool = ToolBuilder::new("add")
2698            .description("Add two numbers")
2699            .handler(|input: AddInput| async move {
2700                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
2701            })
2702            .build();
2703
2704        let router = McpRouter::new().tool(add_tool);
2705        let mut service = JsonRpcService::new(router.clone());
2706
2707        init_jsonrpc_service(&mut service, &router).await;
2708
2709        // Create a batch with one valid and one invalid request
2710        let requests = vec![
2711            // Valid request
2712            JsonRpcRequest::new(1, "tools/call").with_params(serde_json::json!({
2713                "name": "add",
2714                "arguments": {"a": 10, "b": 20}
2715            })),
2716            // Invalid request - tool doesn't exist
2717            JsonRpcRequest::new(2, "tools/call").with_params(serde_json::json!({
2718                "name": "nonexistent_tool",
2719                "arguments": {}
2720            })),
2721            // Another valid request
2722            JsonRpcRequest::new(3, "ping"),
2723        ];
2724
2725        let responses = service.call_batch(requests).await.unwrap();
2726
2727        // All three requests should have responses (errors are not dropped)
2728        assert_eq!(responses.len(), 3);
2729
2730        // First should be success
2731        match &responses[0] {
2732            JsonRpcResponse::Result(r) => {
2733                assert_eq!(r.id, RequestId::Number(1));
2734            }
2735            JsonRpcResponse::Error(_) => panic!("Expected success for first request"),
2736            _ => panic!("unexpected response variant"),
2737        }
2738
2739        // Second should be an error (tool not found)
2740        match &responses[1] {
2741            JsonRpcResponse::Error(e) => {
2742                assert_eq!(e.id, Some(RequestId::Number(2)));
2743                // Error should indicate method not found
2744                assert!(e.error.message.contains("not found") || e.error.code == -32601);
2745            }
2746            JsonRpcResponse::Result(_) => panic!("Expected error for second request"),
2747            _ => panic!("unexpected response variant"),
2748        }
2749
2750        // Third should be success
2751        match &responses[2] {
2752            JsonRpcResponse::Result(r) => {
2753                assert_eq!(r.id, RequestId::Number(3));
2754            }
2755            JsonRpcResponse::Error(_) => panic!("Expected success for third request"),
2756            _ => panic!("unexpected response variant"),
2757        }
2758    }
2759
2760    // =========================================================================
2761    // Resource Template Tests
2762    // =========================================================================
2763
2764    #[tokio::test]
2765    async fn test_list_resource_templates() {
2766        use crate::resource::ResourceTemplateBuilder;
2767        use std::collections::HashMap;
2768
2769        let template = ResourceTemplateBuilder::new("file:///{path}")
2770            .name("Project Files")
2771            .description("Access project files")
2772            .handler(|uri: String, _vars: HashMap<String, String>| async move {
2773                Ok(ReadResourceResult {
2774                    contents: vec![ResourceContent {
2775                        uri,
2776                        mime_type: None,
2777                        text: None,
2778                        blob: None,
2779                        meta: None,
2780                    }],
2781                    meta: None,
2782                })
2783            });
2784
2785        let mut router = McpRouter::new().resource_template(template);
2786
2787        // Initialize session
2788        init_router(&mut router).await;
2789
2790        let req = RouterRequest {
2791            id: RequestId::Number(1),
2792            inner: McpRequest::ListResourceTemplates(ListResourceTemplatesParams::default()),
2793            extensions: Extensions::new(),
2794        };
2795
2796        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2797
2798        match resp.inner {
2799            Ok(McpResponse::ListResourceTemplates(result)) => {
2800                assert_eq!(result.resource_templates.len(), 1);
2801                assert_eq!(result.resource_templates[0].uri_template, "file:///{path}");
2802                assert_eq!(result.resource_templates[0].name, "Project Files");
2803            }
2804            _ => panic!("Expected ListResourceTemplates response"),
2805        }
2806    }
2807
2808    #[tokio::test]
2809    async fn test_read_resource_via_template() {
2810        use crate::resource::ResourceTemplateBuilder;
2811        use std::collections::HashMap;
2812
2813        let template = ResourceTemplateBuilder::new("db://users/{id}")
2814            .name("User Records")
2815            .handler(|uri: String, vars: HashMap<String, String>| async move {
2816                let id = vars.get("id").unwrap().clone();
2817                Ok(ReadResourceResult {
2818                    contents: vec![ResourceContent {
2819                        uri,
2820                        mime_type: Some("application/json".to_string()),
2821                        text: Some(format!(r#"{{"id": "{}"}}"#, id)),
2822                        blob: None,
2823                        meta: None,
2824                    }],
2825                    meta: None,
2826                })
2827            });
2828
2829        let mut router = McpRouter::new().resource_template(template);
2830
2831        // Initialize session
2832        init_router(&mut router).await;
2833
2834        // Read a resource that matches the template
2835        let req = RouterRequest {
2836            id: RequestId::Number(1),
2837            inner: McpRequest::ReadResource(ReadResourceParams {
2838                uri: "db://users/123".to_string(),
2839                meta: None,
2840            }),
2841            extensions: Extensions::new(),
2842        };
2843
2844        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2845
2846        match resp.inner {
2847            Ok(McpResponse::ReadResource(result)) => {
2848                assert_eq!(result.contents.len(), 1);
2849                assert_eq!(result.contents[0].uri, "db://users/123");
2850                assert!(result.contents[0].text.as_ref().unwrap().contains("123"));
2851            }
2852            _ => panic!("Expected ReadResource response"),
2853        }
2854    }
2855
2856    #[tokio::test]
2857    async fn test_static_resource_takes_precedence_over_template() {
2858        use crate::resource::{ResourceBuilder, ResourceTemplateBuilder};
2859        use std::collections::HashMap;
2860
2861        // Template that would match the same URI
2862        let template = ResourceTemplateBuilder::new("file:///{path}")
2863            .name("Files Template")
2864            .handler(|uri: String, _vars: HashMap<String, String>| async move {
2865                Ok(ReadResourceResult {
2866                    contents: vec![ResourceContent {
2867                        uri,
2868                        mime_type: None,
2869                        text: Some("from template".to_string()),
2870                        blob: None,
2871                        meta: None,
2872                    }],
2873                    meta: None,
2874                })
2875            });
2876
2877        // Static resource with exact URI
2878        let static_resource = ResourceBuilder::new("file:///README.md")
2879            .name("README")
2880            .text("from static resource");
2881
2882        let mut router = McpRouter::new()
2883            .resource_template(template)
2884            .resource(static_resource);
2885
2886        // Initialize session
2887        init_router(&mut router).await;
2888
2889        // Read the static resource - should NOT go through template
2890        let req = RouterRequest {
2891            id: RequestId::Number(1),
2892            inner: McpRequest::ReadResource(ReadResourceParams {
2893                uri: "file:///README.md".to_string(),
2894                meta: None,
2895            }),
2896            extensions: Extensions::new(),
2897        };
2898
2899        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2900
2901        match resp.inner {
2902            Ok(McpResponse::ReadResource(result)) => {
2903                // Should get static resource, not template
2904                assert_eq!(
2905                    result.contents[0].text.as_deref(),
2906                    Some("from static resource")
2907                );
2908            }
2909            _ => panic!("Expected ReadResource response"),
2910        }
2911    }
2912
2913    #[tokio::test]
2914    async fn test_resource_not_found_when_no_match() {
2915        use crate::resource::ResourceTemplateBuilder;
2916        use std::collections::HashMap;
2917
2918        let template = ResourceTemplateBuilder::new("db://users/{id}")
2919            .name("Users")
2920            .handler(|uri: String, _vars: HashMap<String, String>| async move {
2921                Ok(ReadResourceResult {
2922                    contents: vec![ResourceContent {
2923                        uri,
2924                        mime_type: None,
2925                        text: None,
2926                        blob: None,
2927                        meta: None,
2928                    }],
2929                    meta: None,
2930                })
2931            });
2932
2933        let mut router = McpRouter::new().resource_template(template);
2934
2935        // Initialize session
2936        init_router(&mut router).await;
2937
2938        // Try to read a URI that doesn't match any resource or template
2939        let req = RouterRequest {
2940            id: RequestId::Number(1),
2941            inner: McpRequest::ReadResource(ReadResourceParams {
2942                uri: "db://posts/123".to_string(),
2943                meta: None,
2944            }),
2945            extensions: Extensions::new(),
2946        };
2947
2948        let resp = router.ready().await.unwrap().call(req).await.unwrap();
2949
2950        match resp.inner {
2951            Err(err) => {
2952                assert!(err.message.contains("not found"));
2953            }
2954            Ok(_) => panic!("Expected error for non-matching URI"),
2955        }
2956    }
2957
2958    #[tokio::test]
2959    async fn test_capabilities_include_resources_with_only_templates() {
2960        use crate::resource::ResourceTemplateBuilder;
2961        use std::collections::HashMap;
2962
2963        let template = ResourceTemplateBuilder::new("file:///{path}")
2964            .name("Files")
2965            .handler(|uri: String, _vars: HashMap<String, String>| async move {
2966                Ok(ReadResourceResult {
2967                    contents: vec![ResourceContent {
2968                        uri,
2969                        mime_type: None,
2970                        text: None,
2971                        blob: None,
2972                        meta: None,
2973                    }],
2974                    meta: None,
2975                })
2976            });
2977
2978        let mut router = McpRouter::new().resource_template(template);
2979
2980        // Send initialize request and check capabilities
2981        let init_req = RouterRequest {
2982            id: RequestId::Number(0),
2983            inner: McpRequest::Initialize(InitializeParams {
2984                protocol_version: "2025-11-25".to_string(),
2985                capabilities: ClientCapabilities {
2986                    roots: None,
2987                    sampling: None,
2988                    elicitation: None,
2989                    tasks: None,
2990                    experimental: None,
2991                    extensions: None,
2992                },
2993                client_info: Implementation {
2994                    name: "test".to_string(),
2995                    version: "1.0".to_string(),
2996                    ..Default::default()
2997                },
2998                meta: None,
2999            }),
3000            extensions: Extensions::new(),
3001        };
3002        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3003
3004        match resp.inner {
3005            Ok(McpResponse::Initialize(result)) => {
3006                // Should have resources capability even though only templates registered
3007                assert!(result.capabilities.resources.is_some());
3008            }
3009            _ => panic!("Expected Initialize response"),
3010        }
3011    }
3012
3013    // =========================================================================
3014    // Logging Notification Tests
3015    // =========================================================================
3016
3017    #[tokio::test]
3018    async fn test_log_sends_notification() {
3019        use crate::context::notification_channel;
3020
3021        let (tx, mut rx) = notification_channel(10);
3022        let router = McpRouter::new().with_notification_sender(tx);
3023
3024        // Send an info log
3025        let sent = router.log_info("Test message");
3026        assert!(sent);
3027
3028        // Should receive the notification
3029        let notification = rx.try_recv().unwrap();
3030        match notification {
3031            ServerNotification::LogMessage(params) => {
3032                assert_eq!(params.level, LogLevel::Info);
3033                let data = params.data;
3034                assert_eq!(
3035                    data.get("message").unwrap().as_str().unwrap(),
3036                    "Test message"
3037                );
3038            }
3039            _ => panic!("Expected LogMessage notification"),
3040        }
3041    }
3042
3043    #[tokio::test]
3044    async fn test_log_with_custom_params() {
3045        use crate::context::notification_channel;
3046
3047        let (tx, mut rx) = notification_channel(10);
3048        let router = McpRouter::new().with_notification_sender(tx);
3049
3050        // Send a custom log message
3051        let params = LoggingMessageParams::new(
3052            LogLevel::Error,
3053            serde_json::json!({
3054                "error": "Connection failed",
3055                "host": "localhost"
3056            }),
3057        )
3058        .with_logger("database");
3059
3060        let sent = router.log(params);
3061        assert!(sent);
3062
3063        let notification = rx.try_recv().unwrap();
3064        match notification {
3065            ServerNotification::LogMessage(params) => {
3066                assert_eq!(params.level, LogLevel::Error);
3067                assert_eq!(params.logger.as_deref(), Some("database"));
3068                let data = params.data;
3069                assert_eq!(
3070                    data.get("error").unwrap().as_str().unwrap(),
3071                    "Connection failed"
3072                );
3073            }
3074            _ => panic!("Expected LogMessage notification"),
3075        }
3076    }
3077
3078    #[tokio::test]
3079    async fn test_log_without_channel_returns_false() {
3080        // Router without notification channel
3081        let router = McpRouter::new();
3082
3083        // Should return false when no channel configured
3084        assert!(!router.log_info("Test"));
3085        assert!(!router.log_warning("Test"));
3086        assert!(!router.log_error("Test"));
3087        assert!(!router.log_debug("Test"));
3088    }
3089
3090    #[tokio::test]
3091    async fn test_logging_capability_with_channel() {
3092        use crate::context::notification_channel;
3093
3094        let (tx, _rx) = notification_channel(10);
3095        let mut router = McpRouter::new().with_notification_sender(tx);
3096
3097        // Initialize and check capabilities
3098        let init_req = RouterRequest {
3099            id: RequestId::Number(0),
3100            inner: McpRequest::Initialize(InitializeParams {
3101                protocol_version: "2025-11-25".to_string(),
3102                capabilities: ClientCapabilities {
3103                    roots: None,
3104                    sampling: None,
3105                    elicitation: None,
3106                    tasks: None,
3107                    experimental: None,
3108                    extensions: None,
3109                },
3110                client_info: Implementation {
3111                    name: "test".to_string(),
3112                    version: "1.0".to_string(),
3113                    ..Default::default()
3114                },
3115                meta: None,
3116            }),
3117            extensions: Extensions::new(),
3118        };
3119        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3120
3121        match resp.inner {
3122            Ok(McpResponse::Initialize(result)) => {
3123                // Should have logging capability when notification channel is set
3124                assert!(result.capabilities.logging.is_some());
3125            }
3126            _ => panic!("Expected Initialize response"),
3127        }
3128    }
3129
3130    #[tokio::test]
3131    async fn test_no_logging_capability_without_channel() {
3132        let mut router = McpRouter::new();
3133
3134        // Initialize and check capabilities
3135        let init_req = RouterRequest {
3136            id: RequestId::Number(0),
3137            inner: McpRequest::Initialize(InitializeParams {
3138                protocol_version: "2025-11-25".to_string(),
3139                capabilities: ClientCapabilities {
3140                    roots: None,
3141                    sampling: None,
3142                    elicitation: None,
3143                    tasks: None,
3144                    experimental: None,
3145                    extensions: None,
3146                },
3147                client_info: Implementation {
3148                    name: "test".to_string(),
3149                    version: "1.0".to_string(),
3150                    ..Default::default()
3151                },
3152                meta: None,
3153            }),
3154            extensions: Extensions::new(),
3155        };
3156        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3157
3158        match resp.inner {
3159            Ok(McpResponse::Initialize(result)) => {
3160                // Should NOT have logging capability without notification channel
3161                assert!(result.capabilities.logging.is_none());
3162            }
3163            _ => panic!("Expected Initialize response"),
3164        }
3165    }
3166
3167    // =========================================================================
3168    // Task Lifecycle Tests
3169    // =========================================================================
3170
3171    #[tokio::test]
3172    async fn test_create_task_via_call_tool() {
3173        let add_tool = ToolBuilder::new("add")
3174            .description("Add two numbers")
3175            .task_support(TaskSupportMode::Optional)
3176            .handler(|input: AddInput| async move {
3177                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3178            })
3179            .build();
3180
3181        let mut router = McpRouter::new().tool(add_tool);
3182        init_router(&mut router).await;
3183
3184        let req = RouterRequest {
3185            id: RequestId::Number(1),
3186            inner: McpRequest::CallTool(CallToolParams {
3187                name: "add".to_string(),
3188                arguments: serde_json::json!({"a": 5, "b": 10}),
3189                meta: None,
3190                task: Some(TaskRequestParams { ttl: None }),
3191            }),
3192            extensions: Extensions::new(),
3193        };
3194
3195        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3196
3197        match resp.inner {
3198            Ok(McpResponse::CreateTask(result)) => {
3199                assert!(result.task.task_id.starts_with("task-"));
3200                assert_eq!(result.task.status, TaskStatus::Working);
3201            }
3202            _ => panic!("Expected CreateTask response"),
3203        }
3204    }
3205
3206    #[tokio::test]
3207    async fn test_list_tasks_empty() {
3208        let mut router = McpRouter::new();
3209        init_router(&mut router).await;
3210
3211        let req = RouterRequest {
3212            id: RequestId::Number(1),
3213            inner: McpRequest::ListTasks(ListTasksParams::default()),
3214            extensions: Extensions::new(),
3215        };
3216
3217        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3218
3219        match resp.inner {
3220            Ok(McpResponse::ListTasks(result)) => {
3221                assert!(result.tasks.is_empty());
3222            }
3223            _ => panic!("Expected ListTasks response"),
3224        }
3225    }
3226
3227    #[tokio::test]
3228    async fn test_task_lifecycle_complete() {
3229        let add_tool = ToolBuilder::new("add")
3230            .description("Add two numbers")
3231            .task_support(TaskSupportMode::Optional)
3232            .handler(|input: AddInput| async move {
3233                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3234            })
3235            .build();
3236
3237        let mut router = McpRouter::new().tool(add_tool);
3238        init_router(&mut router).await;
3239
3240        // Create task via tools/call with task params
3241        let req = RouterRequest {
3242            id: RequestId::Number(1),
3243            inner: McpRequest::CallTool(CallToolParams {
3244                name: "add".to_string(),
3245                arguments: serde_json::json!({"a": 7, "b": 8}),
3246                meta: None,
3247                task: Some(TaskRequestParams { ttl: None }),
3248            }),
3249            extensions: Extensions::new(),
3250        };
3251
3252        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3253        let task_id = match resp.inner {
3254            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3255            _ => panic!("Expected CreateTask response"),
3256        };
3257
3258        // Wait for task to complete
3259        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
3260
3261        // Get task result
3262        let req = RouterRequest {
3263            id: RequestId::Number(2),
3264            inner: McpRequest::GetTaskResult(GetTaskResultParams {
3265                task_id: task_id.clone(),
3266                meta: None,
3267            }),
3268            extensions: Extensions::new(),
3269        };
3270
3271        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3272
3273        match resp.inner {
3274            Ok(McpResponse::GetTaskResult(result)) => {
3275                // Result should have _meta with related-task
3276                assert!(result.meta.is_some());
3277                // Check the result content
3278                match &result.content[0] {
3279                    Content::Text { text, .. } => assert_eq!(text, "15"),
3280                    _ => panic!("Expected text content"),
3281                }
3282            }
3283            _ => panic!("Expected GetTaskResult response"),
3284        }
3285    }
3286
3287    #[tokio::test]
3288    async fn test_task_cancellation() {
3289        // Use a slow tool to test cancellation
3290        let slow_tool = ToolBuilder::new("slow")
3291            .description("Slow tool")
3292            .task_support(TaskSupportMode::Optional)
3293            .handler(|_input: serde_json::Value| async move {
3294                tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
3295                Ok(CallToolResult::text("done"))
3296            })
3297            .build();
3298
3299        let mut router = McpRouter::new().tool(slow_tool);
3300        init_router(&mut router).await;
3301
3302        // Create task
3303        let req = RouterRequest {
3304            id: RequestId::Number(1),
3305            inner: McpRequest::CallTool(CallToolParams {
3306                name: "slow".to_string(),
3307                arguments: serde_json::json!({}),
3308                meta: None,
3309                task: Some(TaskRequestParams { ttl: None }),
3310            }),
3311            extensions: Extensions::new(),
3312        };
3313
3314        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3315        let task_id = match resp.inner {
3316            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3317            _ => panic!("Expected CreateTask response"),
3318        };
3319
3320        // Cancel the task
3321        let req = RouterRequest {
3322            id: RequestId::Number(2),
3323            inner: McpRequest::CancelTask(CancelTaskParams {
3324                task_id: task_id.clone(),
3325                reason: Some("Test cancellation".to_string()),
3326                meta: None,
3327            }),
3328            extensions: Extensions::new(),
3329        };
3330
3331        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3332
3333        match resp.inner {
3334            Ok(McpResponse::CancelTask(task_obj)) => {
3335                assert_eq!(task_obj.status, TaskStatus::Cancelled);
3336            }
3337            _ => panic!("Expected CancelTask response"),
3338        }
3339    }
3340
3341    #[tokio::test]
3342    async fn test_get_task_info() {
3343        let add_tool = ToolBuilder::new("add")
3344            .description("Add two numbers")
3345            .task_support(TaskSupportMode::Optional)
3346            .handler(|input: AddInput| async move {
3347                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3348            })
3349            .build();
3350
3351        let mut router = McpRouter::new().tool(add_tool);
3352        init_router(&mut router).await;
3353
3354        // Create task with TTL
3355        let req = RouterRequest {
3356            id: RequestId::Number(1),
3357            inner: McpRequest::CallTool(CallToolParams {
3358                name: "add".to_string(),
3359                arguments: serde_json::json!({"a": 1, "b": 2}),
3360                meta: None,
3361                task: Some(TaskRequestParams { ttl: Some(600_000) }),
3362            }),
3363            extensions: Extensions::new(),
3364        };
3365
3366        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3367        let task_id = match resp.inner {
3368            Ok(McpResponse::CreateTask(result)) => result.task.task_id,
3369            _ => panic!("Expected CreateTask response"),
3370        };
3371
3372        // Get task info
3373        let req = RouterRequest {
3374            id: RequestId::Number(2),
3375            inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3376                task_id: task_id.clone(),
3377                meta: None,
3378            }),
3379            extensions: Extensions::new(),
3380        };
3381
3382        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3383
3384        match resp.inner {
3385            Ok(McpResponse::GetTaskInfo(info)) => {
3386                assert_eq!(info.task_id, task_id);
3387                assert!(info.created_at.contains('T')); // ISO 8601
3388                assert_eq!(info.ttl, Some(600_000));
3389            }
3390            _ => panic!("Expected GetTaskInfo response"),
3391        }
3392    }
3393
3394    #[tokio::test]
3395    async fn test_task_forbidden_tool_rejects_task_params() {
3396        let tool = ToolBuilder::new("sync_only")
3397            .description("Sync only tool")
3398            .handler(|_input: serde_json::Value| async move { Ok(CallToolResult::text("ok")) })
3399            .build();
3400
3401        let mut router = McpRouter::new().tool(tool);
3402        init_router(&mut router).await;
3403
3404        // Try to create task on a tool with Forbidden task support
3405        let req = RouterRequest {
3406            id: RequestId::Number(1),
3407            inner: McpRequest::CallTool(CallToolParams {
3408                name: "sync_only".to_string(),
3409                arguments: serde_json::json!({}),
3410                meta: None,
3411                task: Some(TaskRequestParams { ttl: None }),
3412            }),
3413            extensions: Extensions::new(),
3414        };
3415
3416        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3417
3418        match resp.inner {
3419            Err(e) => {
3420                assert!(e.message.contains("does not support async tasks"));
3421            }
3422            _ => panic!("Expected error response"),
3423        }
3424    }
3425
3426    #[tokio::test]
3427    async fn test_get_nonexistent_task() {
3428        let mut router = McpRouter::new();
3429        init_router(&mut router).await;
3430
3431        let req = RouterRequest {
3432            id: RequestId::Number(1),
3433            inner: McpRequest::GetTaskInfo(GetTaskInfoParams {
3434                task_id: "task-999".to_string(),
3435                meta: None,
3436            }),
3437            extensions: Extensions::new(),
3438        };
3439
3440        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3441
3442        match resp.inner {
3443            Err(e) => {
3444                assert!(e.message.contains("not found"));
3445            }
3446            _ => panic!("Expected error response"),
3447        }
3448    }
3449
3450    // =========================================================================
3451    // Resource Subscription Tests
3452    // =========================================================================
3453
3454    #[tokio::test]
3455    async fn test_subscribe_to_resource() {
3456        use crate::resource::ResourceBuilder;
3457
3458        let resource = ResourceBuilder::new("file:///test.txt")
3459            .name("Test File")
3460            .text("Hello");
3461
3462        let mut router = McpRouter::new().resource(resource);
3463        init_router(&mut router).await;
3464
3465        // Subscribe to the resource
3466        let req = RouterRequest {
3467            id: RequestId::Number(1),
3468            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3469                uri: "file:///test.txt".to_string(),
3470                meta: None,
3471            }),
3472            extensions: Extensions::new(),
3473        };
3474
3475        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3476
3477        match resp.inner {
3478            Ok(McpResponse::SubscribeResource(_)) => {
3479                // Should be subscribed now
3480                assert!(router.is_subscribed("file:///test.txt"));
3481            }
3482            _ => panic!("Expected SubscribeResource response"),
3483        }
3484    }
3485
3486    #[tokio::test]
3487    async fn test_unsubscribe_from_resource() {
3488        use crate::resource::ResourceBuilder;
3489
3490        let resource = ResourceBuilder::new("file:///test.txt")
3491            .name("Test File")
3492            .text("Hello");
3493
3494        let mut router = McpRouter::new().resource(resource);
3495        init_router(&mut router).await;
3496
3497        // Subscribe first
3498        let req = RouterRequest {
3499            id: RequestId::Number(1),
3500            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3501                uri: "file:///test.txt".to_string(),
3502                meta: None,
3503            }),
3504            extensions: Extensions::new(),
3505        };
3506        let _ = router.ready().await.unwrap().call(req).await.unwrap();
3507        assert!(router.is_subscribed("file:///test.txt"));
3508
3509        // Now unsubscribe
3510        let req = RouterRequest {
3511            id: RequestId::Number(2),
3512            inner: McpRequest::UnsubscribeResource(UnsubscribeResourceParams {
3513                uri: "file:///test.txt".to_string(),
3514                meta: None,
3515            }),
3516            extensions: Extensions::new(),
3517        };
3518
3519        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3520
3521        match resp.inner {
3522            Ok(McpResponse::UnsubscribeResource(_)) => {
3523                // Should no longer be subscribed
3524                assert!(!router.is_subscribed("file:///test.txt"));
3525            }
3526            _ => panic!("Expected UnsubscribeResource response"),
3527        }
3528    }
3529
3530    #[tokio::test]
3531    async fn test_subscribe_nonexistent_resource() {
3532        let mut router = McpRouter::new();
3533        init_router(&mut router).await;
3534
3535        let req = RouterRequest {
3536            id: RequestId::Number(1),
3537            inner: McpRequest::SubscribeResource(SubscribeResourceParams {
3538                uri: "file:///nonexistent.txt".to_string(),
3539                meta: None,
3540            }),
3541            extensions: Extensions::new(),
3542        };
3543
3544        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3545
3546        match resp.inner {
3547            Err(e) => {
3548                assert!(e.message.contains("not found"));
3549            }
3550            _ => panic!("Expected error response"),
3551        }
3552    }
3553
3554    #[tokio::test]
3555    async fn test_notify_resource_updated() {
3556        use crate::context::notification_channel;
3557        use crate::resource::ResourceBuilder;
3558
3559        let (tx, mut rx) = notification_channel(10);
3560
3561        let resource = ResourceBuilder::new("file:///test.txt")
3562            .name("Test File")
3563            .text("Hello");
3564
3565        let router = McpRouter::new()
3566            .resource(resource)
3567            .with_notification_sender(tx);
3568
3569        // First, manually subscribe (simulate subscription)
3570        router.subscribe("file:///test.txt");
3571
3572        // Now notify
3573        let sent = router.notify_resource_updated("file:///test.txt");
3574        assert!(sent);
3575
3576        // Check the notification was sent
3577        let notification = rx.try_recv().unwrap();
3578        match notification {
3579            ServerNotification::ResourceUpdated { uri } => {
3580                assert_eq!(uri, "file:///test.txt");
3581            }
3582            _ => panic!("Expected ResourceUpdated notification"),
3583        }
3584    }
3585
3586    #[tokio::test]
3587    async fn test_notify_resource_updated_not_subscribed() {
3588        use crate::context::notification_channel;
3589        use crate::resource::ResourceBuilder;
3590
3591        let (tx, mut rx) = notification_channel(10);
3592
3593        let resource = ResourceBuilder::new("file:///test.txt")
3594            .name("Test File")
3595            .text("Hello");
3596
3597        let router = McpRouter::new()
3598            .resource(resource)
3599            .with_notification_sender(tx);
3600
3601        // Try to notify without subscribing
3602        let sent = router.notify_resource_updated("file:///test.txt");
3603        assert!(!sent); // Should not send because not subscribed
3604
3605        // Channel should be empty
3606        assert!(rx.try_recv().is_err());
3607    }
3608
3609    #[tokio::test]
3610    async fn test_notify_resources_list_changed() {
3611        use crate::context::notification_channel;
3612
3613        let (tx, mut rx) = notification_channel(10);
3614        let router = McpRouter::new().with_notification_sender(tx);
3615
3616        let sent = router.notify_resources_list_changed();
3617        assert!(sent);
3618
3619        let notification = rx.try_recv().unwrap();
3620        match notification {
3621            ServerNotification::ResourcesListChanged => {}
3622            _ => panic!("Expected ResourcesListChanged notification"),
3623        }
3624    }
3625
3626    #[tokio::test]
3627    async fn test_subscribed_uris() {
3628        use crate::resource::ResourceBuilder;
3629
3630        let resource1 = ResourceBuilder::new("file:///a.txt").name("A").text("A");
3631
3632        let resource2 = ResourceBuilder::new("file:///b.txt").name("B").text("B");
3633
3634        let router = McpRouter::new().resource(resource1).resource(resource2);
3635
3636        // Subscribe to both
3637        router.subscribe("file:///a.txt");
3638        router.subscribe("file:///b.txt");
3639
3640        let uris = router.subscribed_uris();
3641        assert_eq!(uris.len(), 2);
3642        assert!(uris.contains(&"file:///a.txt".to_string()));
3643        assert!(uris.contains(&"file:///b.txt".to_string()));
3644    }
3645
3646    #[tokio::test]
3647    async fn test_subscription_capability_advertised() {
3648        use crate::resource::ResourceBuilder;
3649
3650        let resource = ResourceBuilder::new("file:///test.txt")
3651            .name("Test")
3652            .text("Hello");
3653
3654        let mut router = McpRouter::new().resource(resource);
3655
3656        // Initialize and check capabilities
3657        let init_req = RouterRequest {
3658            id: RequestId::Number(0),
3659            inner: McpRequest::Initialize(InitializeParams {
3660                protocol_version: "2025-11-25".to_string(),
3661                capabilities: ClientCapabilities {
3662                    roots: None,
3663                    sampling: None,
3664                    elicitation: None,
3665                    tasks: None,
3666                    experimental: None,
3667                    extensions: None,
3668                },
3669                client_info: Implementation {
3670                    name: "test".to_string(),
3671                    version: "1.0".to_string(),
3672                    ..Default::default()
3673                },
3674                meta: None,
3675            }),
3676            extensions: Extensions::new(),
3677        };
3678        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
3679
3680        match resp.inner {
3681            Ok(McpResponse::Initialize(result)) => {
3682                // Should have resources capability with subscribe enabled
3683                let resources_cap = result.capabilities.resources.unwrap();
3684                assert!(resources_cap.subscribe);
3685            }
3686            _ => panic!("Expected Initialize response"),
3687        }
3688    }
3689
3690    #[tokio::test]
3691    async fn test_completion_handler() {
3692        let router = McpRouter::new()
3693            .server_info("test", "1.0")
3694            .completion_handler(|params: CompleteParams| async move {
3695                // Return suggestions based on the argument value
3696                let prefix = &params.argument.value;
3697                let suggestions: Vec<String> = vec!["alpha", "beta", "gamma"]
3698                    .into_iter()
3699                    .filter(|s| s.starts_with(prefix))
3700                    .map(String::from)
3701                    .collect();
3702                Ok(CompleteResult::new(suggestions))
3703            });
3704
3705        // Initialize
3706        let init_req = RouterRequest {
3707            id: RequestId::Number(0),
3708            inner: McpRequest::Initialize(InitializeParams {
3709                protocol_version: "2025-11-25".to_string(),
3710                capabilities: ClientCapabilities::default(),
3711                client_info: Implementation {
3712                    name: "test".to_string(),
3713                    version: "1.0".to_string(),
3714                    ..Default::default()
3715                },
3716                meta: None,
3717            }),
3718            extensions: Extensions::new(),
3719        };
3720        let resp = router
3721            .clone()
3722            .ready()
3723            .await
3724            .unwrap()
3725            .call(init_req)
3726            .await
3727            .unwrap();
3728
3729        // Check that completions capability is advertised
3730        match resp.inner {
3731            Ok(McpResponse::Initialize(result)) => {
3732                assert!(result.capabilities.completions.is_some());
3733            }
3734            _ => panic!("Expected Initialize response"),
3735        }
3736
3737        // Send initialized notification
3738        router.handle_notification(McpNotification::Initialized);
3739
3740        // Test completion request
3741        let complete_req = RouterRequest {
3742            id: RequestId::Number(1),
3743            inner: McpRequest::Complete(CompleteParams {
3744                reference: CompletionReference::prompt("test-prompt"),
3745                argument: CompletionArgument::new("query", "al"),
3746                context: None,
3747                meta: None,
3748            }),
3749            extensions: Extensions::new(),
3750        };
3751        let resp = router
3752            .clone()
3753            .ready()
3754            .await
3755            .unwrap()
3756            .call(complete_req)
3757            .await
3758            .unwrap();
3759
3760        match resp.inner {
3761            Ok(McpResponse::Complete(result)) => {
3762                assert_eq!(result.completion.values, vec!["alpha"]);
3763            }
3764            _ => panic!("Expected Complete response"),
3765        }
3766    }
3767
3768    #[tokio::test]
3769    async fn test_completion_without_handler_returns_empty() {
3770        let router = McpRouter::new().server_info("test", "1.0");
3771
3772        // Initialize
3773        let init_req = RouterRequest {
3774            id: RequestId::Number(0),
3775            inner: McpRequest::Initialize(InitializeParams {
3776                protocol_version: "2025-11-25".to_string(),
3777                capabilities: ClientCapabilities::default(),
3778                client_info: Implementation {
3779                    name: "test".to_string(),
3780                    version: "1.0".to_string(),
3781                    ..Default::default()
3782                },
3783                meta: None,
3784            }),
3785            extensions: Extensions::new(),
3786        };
3787        let resp = router
3788            .clone()
3789            .ready()
3790            .await
3791            .unwrap()
3792            .call(init_req)
3793            .await
3794            .unwrap();
3795
3796        // Check that completions capability is NOT advertised
3797        match resp.inner {
3798            Ok(McpResponse::Initialize(result)) => {
3799                assert!(result.capabilities.completions.is_none());
3800            }
3801            _ => panic!("Expected Initialize response"),
3802        }
3803
3804        // Send initialized notification
3805        router.handle_notification(McpNotification::Initialized);
3806
3807        // Test completion request still works but returns empty
3808        let complete_req = RouterRequest {
3809            id: RequestId::Number(1),
3810            inner: McpRequest::Complete(CompleteParams {
3811                reference: CompletionReference::prompt("test-prompt"),
3812                argument: CompletionArgument::new("query", "al"),
3813                context: None,
3814                meta: None,
3815            }),
3816            extensions: Extensions::new(),
3817        };
3818        let resp = router
3819            .clone()
3820            .ready()
3821            .await
3822            .unwrap()
3823            .call(complete_req)
3824            .await
3825            .unwrap();
3826
3827        match resp.inner {
3828            Ok(McpResponse::Complete(result)) => {
3829                assert!(result.completion.values.is_empty());
3830            }
3831            _ => panic!("Expected Complete response"),
3832        }
3833    }
3834
3835    #[tokio::test]
3836    async fn test_tool_filter_list() {
3837        use crate::filter::CapabilityFilter;
3838        use crate::tool::Tool;
3839
3840        let public_tool = ToolBuilder::new("public")
3841            .description("Public tool")
3842            .handler(|_: AddInput| async move { Ok(CallToolResult::text("public")) })
3843            .build();
3844
3845        let admin_tool = ToolBuilder::new("admin")
3846            .description("Admin tool")
3847            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3848            .build();
3849
3850        let mut router = McpRouter::new()
3851            .tool(public_tool)
3852            .tool(admin_tool)
3853            .tool_filter(CapabilityFilter::new(|_, tool: &Tool| tool.name != "admin"));
3854
3855        // Initialize session
3856        init_router(&mut router).await;
3857
3858        let req = RouterRequest {
3859            id: RequestId::Number(1),
3860            inner: McpRequest::ListTools(ListToolsParams::default()),
3861            extensions: Extensions::new(),
3862        };
3863
3864        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3865
3866        match resp.inner {
3867            Ok(McpResponse::ListTools(result)) => {
3868                // Only public tool should be visible
3869                assert_eq!(result.tools.len(), 1);
3870                assert_eq!(result.tools[0].name, "public");
3871            }
3872            _ => panic!("Expected ListTools response"),
3873        }
3874    }
3875
3876    #[tokio::test]
3877    async fn test_tool_filter_call_denied() {
3878        use crate::filter::CapabilityFilter;
3879        use crate::tool::Tool;
3880
3881        let admin_tool = ToolBuilder::new("admin")
3882            .description("Admin tool")
3883            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3884            .build();
3885
3886        let mut router = McpRouter::new()
3887            .tool(admin_tool)
3888            .tool_filter(CapabilityFilter::new(|_, _: &Tool| false)); // Deny all
3889
3890        // Initialize session
3891        init_router(&mut router).await;
3892
3893        let req = RouterRequest {
3894            id: RequestId::Number(1),
3895            inner: McpRequest::CallTool(CallToolParams {
3896                name: "admin".to_string(),
3897                arguments: serde_json::json!({"a": 1, "b": 2}),
3898                meta: None,
3899                task: None,
3900            }),
3901            extensions: Extensions::new(),
3902        };
3903
3904        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3905
3906        // Should get method not found error (default denial behavior)
3907        match resp.inner {
3908            Err(e) => {
3909                assert_eq!(e.code, -32601); // Method not found
3910            }
3911            _ => panic!("Expected JsonRpc error"),
3912        }
3913    }
3914
3915    #[tokio::test]
3916    async fn test_tool_filter_call_allowed() {
3917        use crate::filter::CapabilityFilter;
3918        use crate::tool::Tool;
3919
3920        let public_tool = ToolBuilder::new("public")
3921            .description("Public tool")
3922            .handler(|input: AddInput| async move {
3923                Ok(CallToolResult::text(format!("{}", input.a + input.b)))
3924            })
3925            .build();
3926
3927        let mut router = McpRouter::new()
3928            .tool(public_tool)
3929            .tool_filter(CapabilityFilter::new(|_, _: &Tool| true)); // Allow all
3930
3931        // Initialize session
3932        init_router(&mut router).await;
3933
3934        let req = RouterRequest {
3935            id: RequestId::Number(1),
3936            inner: McpRequest::CallTool(CallToolParams {
3937                name: "public".to_string(),
3938                arguments: serde_json::json!({"a": 1, "b": 2}),
3939                meta: None,
3940                task: None,
3941            }),
3942            extensions: Extensions::new(),
3943        };
3944
3945        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3946
3947        match resp.inner {
3948            Ok(McpResponse::CallTool(result)) => {
3949                assert!(!result.is_error);
3950            }
3951            _ => panic!("Expected CallTool response"),
3952        }
3953    }
3954
3955    #[tokio::test]
3956    async fn test_tool_filter_custom_denial() {
3957        use crate::filter::{CapabilityFilter, DenialBehavior};
3958        use crate::tool::Tool;
3959
3960        let admin_tool = ToolBuilder::new("admin")
3961            .description("Admin tool")
3962            .handler(|_: AddInput| async move { Ok(CallToolResult::text("admin")) })
3963            .build();
3964
3965        let mut router = McpRouter::new().tool(admin_tool).tool_filter(
3966            CapabilityFilter::new(|_, _: &Tool| false)
3967                .denial_behavior(DenialBehavior::Unauthorized),
3968        );
3969
3970        // Initialize session
3971        init_router(&mut router).await;
3972
3973        let req = RouterRequest {
3974            id: RequestId::Number(1),
3975            inner: McpRequest::CallTool(CallToolParams {
3976                name: "admin".to_string(),
3977                arguments: serde_json::json!({"a": 1, "b": 2}),
3978                meta: None,
3979                task: None,
3980            }),
3981            extensions: Extensions::new(),
3982        };
3983
3984        let resp = router.ready().await.unwrap().call(req).await.unwrap();
3985
3986        // Should get forbidden error
3987        match resp.inner {
3988            Err(e) => {
3989                assert_eq!(e.code, -32007); // Forbidden
3990                assert!(e.message.contains("Unauthorized"));
3991            }
3992            _ => panic!("Expected JsonRpc error"),
3993        }
3994    }
3995
3996    #[tokio::test]
3997    async fn test_resource_filter_list() {
3998        use crate::filter::CapabilityFilter;
3999        use crate::resource::{Resource, ResourceBuilder};
4000
4001        let public_resource = ResourceBuilder::new("file:///public.txt")
4002            .name("Public File")
4003            .text("public content");
4004
4005        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4006            .name("Secret File")
4007            .text("secret content");
4008
4009        let mut router = McpRouter::new()
4010            .resource(public_resource)
4011            .resource(secret_resource)
4012            .resource_filter(CapabilityFilter::new(|_, r: &Resource| {
4013                !r.name.contains("Secret")
4014            }));
4015
4016        // Initialize session
4017        init_router(&mut router).await;
4018
4019        let req = RouterRequest {
4020            id: RequestId::Number(1),
4021            inner: McpRequest::ListResources(ListResourcesParams::default()),
4022            extensions: Extensions::new(),
4023        };
4024
4025        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4026
4027        match resp.inner {
4028            Ok(McpResponse::ListResources(result)) => {
4029                // Should only see public resource
4030                assert_eq!(result.resources.len(), 1);
4031                assert_eq!(result.resources[0].name, "Public File");
4032            }
4033            _ => panic!("Expected ListResources response"),
4034        }
4035    }
4036
4037    #[tokio::test]
4038    async fn test_resource_filter_read_denied() {
4039        use crate::filter::CapabilityFilter;
4040        use crate::resource::{Resource, ResourceBuilder};
4041
4042        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4043            .name("Secret File")
4044            .text("secret content");
4045
4046        let mut router = McpRouter::new()
4047            .resource(secret_resource)
4048            .resource_filter(CapabilityFilter::new(|_, _: &Resource| false)); // Deny all
4049
4050        // Initialize session
4051        init_router(&mut router).await;
4052
4053        let req = RouterRequest {
4054            id: RequestId::Number(1),
4055            inner: McpRequest::ReadResource(ReadResourceParams {
4056                uri: "file:///secret.txt".to_string(),
4057                meta: None,
4058            }),
4059            extensions: Extensions::new(),
4060        };
4061
4062        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4063
4064        // Should get method not found error (default denial behavior)
4065        match resp.inner {
4066            Err(e) => {
4067                assert_eq!(e.code, -32601); // Method not found
4068            }
4069            _ => panic!("Expected JsonRpc error"),
4070        }
4071    }
4072
4073    #[tokio::test]
4074    async fn test_resource_filter_read_allowed() {
4075        use crate::filter::CapabilityFilter;
4076        use crate::resource::{Resource, ResourceBuilder};
4077
4078        let public_resource = ResourceBuilder::new("file:///public.txt")
4079            .name("Public File")
4080            .text("public content");
4081
4082        let mut router = McpRouter::new()
4083            .resource(public_resource)
4084            .resource_filter(CapabilityFilter::new(|_, _: &Resource| true)); // Allow all
4085
4086        // Initialize session
4087        init_router(&mut router).await;
4088
4089        let req = RouterRequest {
4090            id: RequestId::Number(1),
4091            inner: McpRequest::ReadResource(ReadResourceParams {
4092                uri: "file:///public.txt".to_string(),
4093                meta: None,
4094            }),
4095            extensions: Extensions::new(),
4096        };
4097
4098        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4099
4100        match resp.inner {
4101            Ok(McpResponse::ReadResource(result)) => {
4102                assert_eq!(result.contents.len(), 1);
4103                assert_eq!(result.contents[0].text.as_deref(), Some("public content"));
4104            }
4105            _ => panic!("Expected ReadResource response"),
4106        }
4107    }
4108
4109    #[tokio::test]
4110    async fn test_resource_filter_custom_denial() {
4111        use crate::filter::{CapabilityFilter, DenialBehavior};
4112        use crate::resource::{Resource, ResourceBuilder};
4113
4114        let secret_resource = ResourceBuilder::new("file:///secret.txt")
4115            .name("Secret File")
4116            .text("secret content");
4117
4118        let mut router = McpRouter::new().resource(secret_resource).resource_filter(
4119            CapabilityFilter::new(|_, _: &Resource| false)
4120                .denial_behavior(DenialBehavior::Unauthorized),
4121        );
4122
4123        // Initialize session
4124        init_router(&mut router).await;
4125
4126        let req = RouterRequest {
4127            id: RequestId::Number(1),
4128            inner: McpRequest::ReadResource(ReadResourceParams {
4129                uri: "file:///secret.txt".to_string(),
4130                meta: None,
4131            }),
4132            extensions: Extensions::new(),
4133        };
4134
4135        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4136
4137        // Should get forbidden error
4138        match resp.inner {
4139            Err(e) => {
4140                assert_eq!(e.code, -32007); // Forbidden
4141                assert!(e.message.contains("Unauthorized"));
4142            }
4143            _ => panic!("Expected JsonRpc error"),
4144        }
4145    }
4146
4147    #[tokio::test]
4148    async fn test_prompt_filter_list() {
4149        use crate::filter::CapabilityFilter;
4150        use crate::prompt::{Prompt, PromptBuilder};
4151
4152        let public_prompt = PromptBuilder::new("greeting")
4153            .description("A greeting")
4154            .user_message("Hello!");
4155
4156        let admin_prompt = PromptBuilder::new("system_debug")
4157            .description("Admin prompt")
4158            .user_message("Debug");
4159
4160        let mut router = McpRouter::new()
4161            .prompt(public_prompt)
4162            .prompt(admin_prompt)
4163            .prompt_filter(CapabilityFilter::new(|_, p: &Prompt| {
4164                !p.name.contains("system")
4165            }));
4166
4167        // Initialize session
4168        init_router(&mut router).await;
4169
4170        let req = RouterRequest {
4171            id: RequestId::Number(1),
4172            inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4173            extensions: Extensions::new(),
4174        };
4175
4176        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4177
4178        match resp.inner {
4179            Ok(McpResponse::ListPrompts(result)) => {
4180                // Should only see public prompt
4181                assert_eq!(result.prompts.len(), 1);
4182                assert_eq!(result.prompts[0].name, "greeting");
4183            }
4184            _ => panic!("Expected ListPrompts response"),
4185        }
4186    }
4187
4188    #[tokio::test]
4189    async fn test_prompt_filter_get_denied() {
4190        use crate::filter::CapabilityFilter;
4191        use crate::prompt::{Prompt, PromptBuilder};
4192        use std::collections::HashMap;
4193
4194        let admin_prompt = PromptBuilder::new("system_debug")
4195            .description("Admin prompt")
4196            .user_message("Debug");
4197
4198        let mut router = McpRouter::new()
4199            .prompt(admin_prompt)
4200            .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| false)); // Deny all
4201
4202        // Initialize session
4203        init_router(&mut router).await;
4204
4205        let req = RouterRequest {
4206            id: RequestId::Number(1),
4207            inner: McpRequest::GetPrompt(GetPromptParams {
4208                name: "system_debug".to_string(),
4209                arguments: HashMap::new(),
4210                meta: None,
4211            }),
4212            extensions: Extensions::new(),
4213        };
4214
4215        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4216
4217        // Should get method not found error (default denial behavior)
4218        match resp.inner {
4219            Err(e) => {
4220                assert_eq!(e.code, -32601); // Method not found
4221            }
4222            _ => panic!("Expected JsonRpc error"),
4223        }
4224    }
4225
4226    #[tokio::test]
4227    async fn test_prompt_filter_get_allowed() {
4228        use crate::filter::CapabilityFilter;
4229        use crate::prompt::{Prompt, PromptBuilder};
4230        use std::collections::HashMap;
4231
4232        let public_prompt = PromptBuilder::new("greeting")
4233            .description("A greeting")
4234            .user_message("Hello!");
4235
4236        let mut router = McpRouter::new()
4237            .prompt(public_prompt)
4238            .prompt_filter(CapabilityFilter::new(|_, _: &Prompt| true)); // Allow all
4239
4240        // Initialize session
4241        init_router(&mut router).await;
4242
4243        let req = RouterRequest {
4244            id: RequestId::Number(1),
4245            inner: McpRequest::GetPrompt(GetPromptParams {
4246                name: "greeting".to_string(),
4247                arguments: HashMap::new(),
4248                meta: None,
4249            }),
4250            extensions: Extensions::new(),
4251        };
4252
4253        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4254
4255        match resp.inner {
4256            Ok(McpResponse::GetPrompt(result)) => {
4257                assert_eq!(result.messages.len(), 1);
4258            }
4259            _ => panic!("Expected GetPrompt response"),
4260        }
4261    }
4262
4263    #[tokio::test]
4264    async fn test_prompt_filter_custom_denial() {
4265        use crate::filter::{CapabilityFilter, DenialBehavior};
4266        use crate::prompt::{Prompt, PromptBuilder};
4267        use std::collections::HashMap;
4268
4269        let admin_prompt = PromptBuilder::new("system_debug")
4270            .description("Admin prompt")
4271            .user_message("Debug");
4272
4273        let mut router = McpRouter::new().prompt(admin_prompt).prompt_filter(
4274            CapabilityFilter::new(|_, _: &Prompt| false)
4275                .denial_behavior(DenialBehavior::Unauthorized),
4276        );
4277
4278        // Initialize session
4279        init_router(&mut router).await;
4280
4281        let req = RouterRequest {
4282            id: RequestId::Number(1),
4283            inner: McpRequest::GetPrompt(GetPromptParams {
4284                name: "system_debug".to_string(),
4285                arguments: HashMap::new(),
4286                meta: None,
4287            }),
4288            extensions: Extensions::new(),
4289        };
4290
4291        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4292
4293        // Should get forbidden error
4294        match resp.inner {
4295            Err(e) => {
4296                assert_eq!(e.code, -32007); // Forbidden
4297                assert!(e.message.contains("Unauthorized"));
4298            }
4299            _ => panic!("Expected JsonRpc error"),
4300        }
4301    }
4302
4303    // =========================================================================
4304    // Router Composition Tests (merge/nest)
4305    // =========================================================================
4306
4307    #[derive(Debug, Deserialize, JsonSchema)]
4308    struct StringInput {
4309        value: String,
4310    }
4311
4312    #[tokio::test]
4313    async fn test_router_merge_tools() {
4314        // Create first router with a tool
4315        let tool_a = ToolBuilder::new("tool_a")
4316            .description("Tool A")
4317            .handler(|_: StringInput| async move { Ok(CallToolResult::text("A")) })
4318            .build();
4319
4320        let router_a = McpRouter::new().tool(tool_a);
4321
4322        // Create second router with different tools
4323        let tool_b = ToolBuilder::new("tool_b")
4324            .description("Tool B")
4325            .handler(|_: StringInput| async move { Ok(CallToolResult::text("B")) })
4326            .build();
4327        let tool_c = ToolBuilder::new("tool_c")
4328            .description("Tool C")
4329            .handler(|_: StringInput| async move { Ok(CallToolResult::text("C")) })
4330            .build();
4331
4332        let router_b = McpRouter::new().tool(tool_b).tool(tool_c);
4333
4334        // Merge them
4335        let mut merged = McpRouter::new()
4336            .server_info("merged", "1.0")
4337            .merge(router_a)
4338            .merge(router_b);
4339
4340        init_router(&mut merged).await;
4341
4342        // List tools
4343        let req = RouterRequest {
4344            id: RequestId::Number(1),
4345            inner: McpRequest::ListTools(ListToolsParams::default()),
4346            extensions: Extensions::new(),
4347        };
4348
4349        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4350
4351        match resp.inner {
4352            Ok(McpResponse::ListTools(result)) => {
4353                assert_eq!(result.tools.len(), 3);
4354                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4355                assert!(names.contains(&"tool_a"));
4356                assert!(names.contains(&"tool_b"));
4357                assert!(names.contains(&"tool_c"));
4358            }
4359            _ => panic!("Expected ListTools response"),
4360        }
4361    }
4362
4363    #[tokio::test]
4364    async fn test_router_merge_overwrites_duplicates() {
4365        // Create first router with a tool
4366        let tool_v1 = ToolBuilder::new("shared")
4367            .description("Version 1")
4368            .handler(|_: StringInput| async move { Ok(CallToolResult::text("v1")) })
4369            .build();
4370
4371        let router_a = McpRouter::new().tool(tool_v1);
4372
4373        // Create second router with same tool name but different description
4374        let tool_v2 = ToolBuilder::new("shared")
4375            .description("Version 2")
4376            .handler(|_: StringInput| async move { Ok(CallToolResult::text("v2")) })
4377            .build();
4378
4379        let router_b = McpRouter::new().tool(tool_v2);
4380
4381        // Merge - second should win
4382        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4383
4384        init_router(&mut merged).await;
4385
4386        let req = RouterRequest {
4387            id: RequestId::Number(1),
4388            inner: McpRequest::ListTools(ListToolsParams::default()),
4389            extensions: Extensions::new(),
4390        };
4391
4392        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4393
4394        match resp.inner {
4395            Ok(McpResponse::ListTools(result)) => {
4396                assert_eq!(result.tools.len(), 1);
4397                assert_eq!(result.tools[0].name, "shared");
4398                assert_eq!(result.tools[0].description.as_deref(), Some("Version 2"));
4399            }
4400            _ => panic!("Expected ListTools response"),
4401        }
4402    }
4403
4404    #[tokio::test]
4405    async fn test_router_merge_resources() {
4406        use crate::resource::ResourceBuilder;
4407
4408        // Create routers with different resources
4409        let router_a = McpRouter::new().resource(
4410            ResourceBuilder::new("file:///a.txt")
4411                .name("File A")
4412                .text("content a"),
4413        );
4414
4415        let router_b = McpRouter::new().resource(
4416            ResourceBuilder::new("file:///b.txt")
4417                .name("File B")
4418                .text("content b"),
4419        );
4420
4421        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4422
4423        init_router(&mut merged).await;
4424
4425        let req = RouterRequest {
4426            id: RequestId::Number(1),
4427            inner: McpRequest::ListResources(ListResourcesParams::default()),
4428            extensions: Extensions::new(),
4429        };
4430
4431        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4432
4433        match resp.inner {
4434            Ok(McpResponse::ListResources(result)) => {
4435                assert_eq!(result.resources.len(), 2);
4436                let uris: Vec<&str> = result.resources.iter().map(|r| r.uri.as_str()).collect();
4437                assert!(uris.contains(&"file:///a.txt"));
4438                assert!(uris.contains(&"file:///b.txt"));
4439            }
4440            _ => panic!("Expected ListResources response"),
4441        }
4442    }
4443
4444    #[tokio::test]
4445    async fn test_router_merge_prompts() {
4446        use crate::prompt::PromptBuilder;
4447
4448        let router_a =
4449            McpRouter::new().prompt(PromptBuilder::new("prompt_a").user_message("Hello A"));
4450
4451        let router_b =
4452            McpRouter::new().prompt(PromptBuilder::new("prompt_b").user_message("Hello B"));
4453
4454        let mut merged = McpRouter::new().merge(router_a).merge(router_b);
4455
4456        init_router(&mut merged).await;
4457
4458        let req = RouterRequest {
4459            id: RequestId::Number(1),
4460            inner: McpRequest::ListPrompts(ListPromptsParams::default()),
4461            extensions: Extensions::new(),
4462        };
4463
4464        let resp = merged.ready().await.unwrap().call(req).await.unwrap();
4465
4466        match resp.inner {
4467            Ok(McpResponse::ListPrompts(result)) => {
4468                assert_eq!(result.prompts.len(), 2);
4469                let names: Vec<&str> = result.prompts.iter().map(|p| p.name.as_str()).collect();
4470                assert!(names.contains(&"prompt_a"));
4471                assert!(names.contains(&"prompt_b"));
4472            }
4473            _ => panic!("Expected ListPrompts response"),
4474        }
4475    }
4476
4477    #[tokio::test]
4478    async fn test_router_nest_prefixes_tools() {
4479        // Create a router with tools
4480        let tool_query = ToolBuilder::new("query")
4481            .description("Query the database")
4482            .handler(|_: StringInput| async move { Ok(CallToolResult::text("query result")) })
4483            .build();
4484        let tool_insert = ToolBuilder::new("insert")
4485            .description("Insert into database")
4486            .handler(|_: StringInput| async move { Ok(CallToolResult::text("insert result")) })
4487            .build();
4488
4489        let db_router = McpRouter::new().tool(tool_query).tool(tool_insert);
4490
4491        // Nest under "db" prefix
4492        let mut router = McpRouter::new()
4493            .server_info("nested", "1.0")
4494            .nest("db", db_router);
4495
4496        init_router(&mut router).await;
4497
4498        let req = RouterRequest {
4499            id: RequestId::Number(1),
4500            inner: McpRequest::ListTools(ListToolsParams::default()),
4501            extensions: Extensions::new(),
4502        };
4503
4504        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4505
4506        match resp.inner {
4507            Ok(McpResponse::ListTools(result)) => {
4508                assert_eq!(result.tools.len(), 2);
4509                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4510                assert!(names.contains(&"db.query"));
4511                assert!(names.contains(&"db.insert"));
4512            }
4513            _ => panic!("Expected ListTools response"),
4514        }
4515    }
4516
4517    #[tokio::test]
4518    async fn test_router_nest_call_prefixed_tool() {
4519        let tool = ToolBuilder::new("echo")
4520            .description("Echo input")
4521            .handler(|input: StringInput| async move { Ok(CallToolResult::text(&input.value)) })
4522            .build();
4523
4524        let nested_router = McpRouter::new().tool(tool);
4525
4526        let mut router = McpRouter::new().nest("api", nested_router);
4527
4528        init_router(&mut router).await;
4529
4530        // Call the prefixed tool
4531        let req = RouterRequest {
4532            id: RequestId::Number(1),
4533            inner: McpRequest::CallTool(CallToolParams {
4534                name: "api.echo".to_string(),
4535                arguments: serde_json::json!({"value": "hello world"}),
4536                meta: None,
4537                task: None,
4538            }),
4539            extensions: Extensions::new(),
4540        };
4541
4542        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4543
4544        match resp.inner {
4545            Ok(McpResponse::CallTool(result)) => {
4546                assert!(!result.is_error);
4547                match &result.content[0] {
4548                    Content::Text { text, .. } => assert_eq!(text, "hello world"),
4549                    _ => panic!("Expected text content"),
4550                }
4551            }
4552            _ => panic!("Expected CallTool response"),
4553        }
4554    }
4555
4556    #[tokio::test]
4557    async fn test_router_multiple_nests() {
4558        let db_tool = ToolBuilder::new("query")
4559            .description("Database query")
4560            .handler(|_: StringInput| async move { Ok(CallToolResult::text("db")) })
4561            .build();
4562
4563        let api_tool = ToolBuilder::new("fetch")
4564            .description("API fetch")
4565            .handler(|_: StringInput| async move { Ok(CallToolResult::text("api")) })
4566            .build();
4567
4568        let db_router = McpRouter::new().tool(db_tool);
4569        let api_router = McpRouter::new().tool(api_tool);
4570
4571        let mut router = McpRouter::new()
4572            .nest("db", db_router)
4573            .nest("api", api_router);
4574
4575        init_router(&mut router).await;
4576
4577        let req = RouterRequest {
4578            id: RequestId::Number(1),
4579            inner: McpRequest::ListTools(ListToolsParams::default()),
4580            extensions: Extensions::new(),
4581        };
4582
4583        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4584
4585        match resp.inner {
4586            Ok(McpResponse::ListTools(result)) => {
4587                assert_eq!(result.tools.len(), 2);
4588                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4589                assert!(names.contains(&"db.query"));
4590                assert!(names.contains(&"api.fetch"));
4591            }
4592            _ => panic!("Expected ListTools response"),
4593        }
4594    }
4595
4596    #[tokio::test]
4597    async fn test_router_merge_and_nest_combined() {
4598        // Test combining merge and nest
4599        let tool_a = ToolBuilder::new("local")
4600            .description("Local tool")
4601            .handler(|_: StringInput| async move { Ok(CallToolResult::text("local")) })
4602            .build();
4603
4604        let nested_tool = ToolBuilder::new("remote")
4605            .description("Remote tool")
4606            .handler(|_: StringInput| async move { Ok(CallToolResult::text("remote")) })
4607            .build();
4608
4609        let nested_router = McpRouter::new().tool(nested_tool);
4610
4611        let mut router = McpRouter::new()
4612            .tool(tool_a)
4613            .nest("external", nested_router);
4614
4615        init_router(&mut router).await;
4616
4617        let req = RouterRequest {
4618            id: RequestId::Number(1),
4619            inner: McpRequest::ListTools(ListToolsParams::default()),
4620            extensions: Extensions::new(),
4621        };
4622
4623        let resp = router.ready().await.unwrap().call(req).await.unwrap();
4624
4625        match resp.inner {
4626            Ok(McpResponse::ListTools(result)) => {
4627                assert_eq!(result.tools.len(), 2);
4628                let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
4629                assert!(names.contains(&"local"));
4630                assert!(names.contains(&"external.remote"));
4631            }
4632            _ => panic!("Expected ListTools response"),
4633        }
4634    }
4635
4636    #[tokio::test]
4637    async fn test_router_merge_preserves_server_info() {
4638        let child_router = McpRouter::new()
4639            .server_info("child", "2.0")
4640            .instructions("Child instructions");
4641
4642        let mut router = McpRouter::new()
4643            .server_info("parent", "1.0")
4644            .instructions("Parent instructions")
4645            .merge(child_router);
4646
4647        init_router(&mut router).await;
4648
4649        // Initialize response should have parent's server info
4650        let init_req = RouterRequest {
4651            id: RequestId::Number(99),
4652            inner: McpRequest::Initialize(InitializeParams {
4653                protocol_version: "2025-11-25".to_string(),
4654                capabilities: ClientCapabilities::default(),
4655                client_info: Implementation {
4656                    name: "test".to_string(),
4657                    version: "1.0".to_string(),
4658                    ..Default::default()
4659                },
4660                meta: None,
4661            }),
4662            extensions: Extensions::new(),
4663        };
4664
4665        // Create fresh router for this test since we need to call initialize
4666        let child_router2 = McpRouter::new().server_info("child", "2.0");
4667        let mut fresh_router = McpRouter::new()
4668            .server_info("parent", "1.0")
4669            .merge(child_router2);
4670
4671        let resp = fresh_router
4672            .ready()
4673            .await
4674            .unwrap()
4675            .call(init_req)
4676            .await
4677            .unwrap();
4678
4679        match resp.inner {
4680            Ok(McpResponse::Initialize(result)) => {
4681                assert_eq!(result.server_info.name, "parent");
4682                assert_eq!(result.server_info.version, "1.0");
4683            }
4684            _ => panic!("Expected Initialize response"),
4685        }
4686    }
4687
4688    // =========================================================================
4689    // Auto-instructions tests
4690    // =========================================================================
4691
4692    #[tokio::test]
4693    async fn test_auto_instructions_tools_only() {
4694        let tool_a = ToolBuilder::new("alpha")
4695            .description("Alpha tool")
4696            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4697            .build();
4698        let tool_b = ToolBuilder::new("beta")
4699            .description("Beta tool")
4700            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4701            .build();
4702
4703        let mut router = McpRouter::new()
4704            .auto_instructions()
4705            .tool(tool_a)
4706            .tool(tool_b);
4707
4708        let resp = send_initialize(&mut router).await;
4709        let instructions = resp.instructions.expect("should have instructions");
4710
4711        assert!(instructions.contains("## Tools"));
4712        assert!(instructions.contains("- **alpha**: Alpha tool"));
4713        assert!(instructions.contains("- **beta**: Beta tool"));
4714        // No resources or prompts sections
4715        assert!(!instructions.contains("## Resources"));
4716        assert!(!instructions.contains("## Prompts"));
4717    }
4718
4719    #[tokio::test]
4720    async fn test_auto_instructions_with_annotations() {
4721        let read_only_tool = ToolBuilder::new("query")
4722            .description("Run a query")
4723            .read_only()
4724            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4725            .build();
4726        let destructive_tool = ToolBuilder::new("delete")
4727            .description("Delete a record")
4728            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4729            .build();
4730        let idempotent_tool = ToolBuilder::new("upsert")
4731            .description("Upsert a record")
4732            .non_destructive()
4733            .idempotent()
4734            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4735            .build();
4736
4737        let mut router = McpRouter::new()
4738            .auto_instructions()
4739            .tool(read_only_tool)
4740            .tool(destructive_tool)
4741            .tool(idempotent_tool);
4742
4743        let resp = send_initialize(&mut router).await;
4744        let instructions = resp.instructions.unwrap();
4745
4746        assert!(instructions.contains("- **query**: Run a query [read-only]"));
4747        // delete has no annotations set via builder, so no tags
4748        assert!(instructions.contains("- **delete**: Delete a record\n"));
4749        assert!(instructions.contains("- **upsert**: Upsert a record [idempotent]"));
4750    }
4751
4752    #[tokio::test]
4753    async fn test_auto_instructions_with_resources() {
4754        use crate::resource::ResourceBuilder;
4755
4756        let resource = ResourceBuilder::new("file:///schema.sql")
4757            .name("Schema")
4758            .description("Database schema")
4759            .text("CREATE TABLE ...");
4760
4761        let mut router = McpRouter::new().auto_instructions().resource(resource);
4762
4763        let resp = send_initialize(&mut router).await;
4764        let instructions = resp.instructions.unwrap();
4765
4766        assert!(instructions.contains("## Resources"));
4767        assert!(instructions.contains("- **file:///schema.sql**: Database schema"));
4768        assert!(!instructions.contains("## Tools"));
4769    }
4770
4771    #[tokio::test]
4772    async fn test_auto_instructions_with_resource_templates() {
4773        use crate::resource::ResourceTemplateBuilder;
4774
4775        let template = ResourceTemplateBuilder::new("file:///{path}")
4776            .name("File")
4777            .description("Read a file by path")
4778            .handler(
4779                |_uri: String, _vars: std::collections::HashMap<String, String>| async move {
4780                    Ok(crate::ReadResourceResult::text("content", "text/plain"))
4781                },
4782            );
4783
4784        let mut router = McpRouter::new()
4785            .auto_instructions()
4786            .resource_template(template);
4787
4788        let resp = send_initialize(&mut router).await;
4789        let instructions = resp.instructions.unwrap();
4790
4791        assert!(instructions.contains("## Resources"));
4792        assert!(instructions.contains("- **file:///{path}**: Read a file by path"));
4793    }
4794
4795    #[tokio::test]
4796    async fn test_auto_instructions_with_prompts() {
4797        use crate::prompt::PromptBuilder;
4798
4799        let prompt = PromptBuilder::new("write_query")
4800            .description("Help write a SQL query")
4801            .user_message("Write a query for: {task}");
4802
4803        let mut router = McpRouter::new().auto_instructions().prompt(prompt);
4804
4805        let resp = send_initialize(&mut router).await;
4806        let instructions = resp.instructions.unwrap();
4807
4808        assert!(instructions.contains("## Prompts"));
4809        assert!(instructions.contains("- **write_query**: Help write a SQL query"));
4810        assert!(!instructions.contains("## Tools"));
4811    }
4812
4813    #[tokio::test]
4814    async fn test_auto_instructions_all_sections() {
4815        use crate::prompt::PromptBuilder;
4816        use crate::resource::ResourceBuilder;
4817
4818        let tool = ToolBuilder::new("query")
4819            .description("Execute SQL")
4820            .read_only()
4821            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4822            .build();
4823        let resource = ResourceBuilder::new("db://schema")
4824            .name("Schema")
4825            .description("Full database schema")
4826            .text("schema");
4827        let prompt = PromptBuilder::new("write_query")
4828            .description("Help write a SQL query")
4829            .user_message("Write a query");
4830
4831        let mut router = McpRouter::new()
4832            .auto_instructions()
4833            .tool(tool)
4834            .resource(resource)
4835            .prompt(prompt);
4836
4837        let resp = send_initialize(&mut router).await;
4838        let instructions = resp.instructions.unwrap();
4839
4840        // All three sections present
4841        assert!(instructions.contains("## Tools"));
4842        assert!(instructions.contains("## Resources"));
4843        assert!(instructions.contains("## Prompts"));
4844
4845        // Sections appear in order: Tools, Resources, Prompts
4846        let tools_pos = instructions.find("## Tools").unwrap();
4847        let resources_pos = instructions.find("## Resources").unwrap();
4848        let prompts_pos = instructions.find("## Prompts").unwrap();
4849        assert!(tools_pos < resources_pos);
4850        assert!(resources_pos < prompts_pos);
4851    }
4852
4853    #[tokio::test]
4854    async fn test_auto_instructions_with_prefix_and_suffix() {
4855        let tool = ToolBuilder::new("echo")
4856            .description("Echo input")
4857            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4858            .build();
4859
4860        let mut router = McpRouter::new()
4861            .auto_instructions_with(
4862                Some("This server provides echo capabilities."),
4863                Some("Contact admin@example.com for support."),
4864            )
4865            .tool(tool);
4866
4867        let resp = send_initialize(&mut router).await;
4868        let instructions = resp.instructions.unwrap();
4869
4870        assert!(instructions.starts_with("This server provides echo capabilities."));
4871        assert!(instructions.ends_with("Contact admin@example.com for support."));
4872        assert!(instructions.contains("## Tools"));
4873        assert!(instructions.contains("- **echo**: Echo input"));
4874    }
4875
4876    #[tokio::test]
4877    async fn test_auto_instructions_prefix_only() {
4878        let tool = ToolBuilder::new("echo")
4879            .description("Echo input")
4880            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4881            .build();
4882
4883        let mut router = McpRouter::new()
4884            .auto_instructions_with(Some("My server intro."), None::<String>)
4885            .tool(tool);
4886
4887        let resp = send_initialize(&mut router).await;
4888        let instructions = resp.instructions.unwrap();
4889
4890        assert!(instructions.starts_with("My server intro."));
4891        assert!(instructions.contains("- **echo**: Echo input"));
4892    }
4893
4894    #[tokio::test]
4895    async fn test_auto_instructions_empty_router() {
4896        let mut router = McpRouter::new().auto_instructions();
4897
4898        let resp = send_initialize(&mut router).await;
4899        let instructions = resp.instructions.expect("should have instructions");
4900
4901        // No sections when nothing is registered
4902        assert!(!instructions.contains("## Tools"));
4903        assert!(!instructions.contains("## Resources"));
4904        assert!(!instructions.contains("## Prompts"));
4905        assert!(instructions.is_empty());
4906    }
4907
4908    #[tokio::test]
4909    async fn test_auto_instructions_overrides_manual() {
4910        let tool = ToolBuilder::new("echo")
4911            .description("Echo input")
4912            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4913            .build();
4914
4915        let mut router = McpRouter::new()
4916            .instructions("This will be overridden")
4917            .auto_instructions()
4918            .tool(tool);
4919
4920        let resp = send_initialize(&mut router).await;
4921        let instructions = resp.instructions.unwrap();
4922
4923        assert!(!instructions.contains("This will be overridden"));
4924        assert!(instructions.contains("- **echo**: Echo input"));
4925    }
4926
4927    #[tokio::test]
4928    async fn test_no_auto_instructions_returns_manual() {
4929        let tool = ToolBuilder::new("echo")
4930            .description("Echo input")
4931            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4932            .build();
4933
4934        let mut router = McpRouter::new()
4935            .instructions("Manual instructions here")
4936            .tool(tool);
4937
4938        let resp = send_initialize(&mut router).await;
4939        let instructions = resp.instructions.unwrap();
4940
4941        assert_eq!(instructions, "Manual instructions here");
4942    }
4943
4944    #[tokio::test]
4945    async fn test_auto_instructions_no_description_fallback() {
4946        let tool = ToolBuilder::new("mystery")
4947            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4948            .build();
4949
4950        let mut router = McpRouter::new().auto_instructions().tool(tool);
4951
4952        let resp = send_initialize(&mut router).await;
4953        let instructions = resp.instructions.unwrap();
4954
4955        assert!(instructions.contains("- **mystery**: No description"));
4956    }
4957
4958    #[tokio::test]
4959    async fn test_auto_instructions_sorted_alphabetically() {
4960        let tool_z = ToolBuilder::new("zebra")
4961            .description("Z tool")
4962            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4963            .build();
4964        let tool_a = ToolBuilder::new("alpha")
4965            .description("A tool")
4966            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4967            .build();
4968        let tool_m = ToolBuilder::new("middle")
4969            .description("M tool")
4970            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4971            .build();
4972
4973        let mut router = McpRouter::new()
4974            .auto_instructions()
4975            .tool(tool_z)
4976            .tool(tool_a)
4977            .tool(tool_m);
4978
4979        let resp = send_initialize(&mut router).await;
4980        let instructions = resp.instructions.unwrap();
4981
4982        let alpha_pos = instructions.find("**alpha**").unwrap();
4983        let middle_pos = instructions.find("**middle**").unwrap();
4984        let zebra_pos = instructions.find("**zebra**").unwrap();
4985        assert!(alpha_pos < middle_pos);
4986        assert!(middle_pos < zebra_pos);
4987    }
4988
4989    #[tokio::test]
4990    async fn test_auto_instructions_read_only_and_idempotent_tags() {
4991        let tool = ToolBuilder::new("safe_update")
4992            .description("Safe update operation")
4993            .idempotent()
4994            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
4995            .build();
4996
4997        let mut router = McpRouter::new().auto_instructions().tool(tool);
4998
4999        let resp = send_initialize(&mut router).await;
5000        let instructions = resp.instructions.unwrap();
5001
5002        assert!(
5003            instructions.contains("[idempotent]"),
5004            "got: {}",
5005            instructions
5006        );
5007    }
5008
5009    #[tokio::test]
5010    async fn test_auto_instructions_lazy_generation() {
5011        // auto_instructions() is called BEFORE tools are registered
5012        // but instructions should still include tools
5013        let mut router = McpRouter::new().auto_instructions();
5014
5015        let tool = ToolBuilder::new("late_tool")
5016            .description("Added after auto_instructions")
5017            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5018            .build();
5019
5020        router = router.tool(tool);
5021
5022        let resp = send_initialize(&mut router).await;
5023        let instructions = resp.instructions.unwrap();
5024
5025        assert!(instructions.contains("- **late_tool**: Added after auto_instructions"));
5026    }
5027
5028    #[tokio::test]
5029    async fn test_auto_instructions_multiple_annotation_tags() {
5030        let tool = ToolBuilder::new("update")
5031            .description("Update a record")
5032            .annotations(ToolAnnotations {
5033                read_only_hint: true,
5034                idempotent_hint: true,
5035                ..Default::default()
5036            })
5037            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5038            .build();
5039
5040        let mut router = McpRouter::new().auto_instructions().tool(tool);
5041
5042        let resp = send_initialize(&mut router).await;
5043        let instructions = resp.instructions.unwrap();
5044
5045        assert!(
5046            instructions.contains("[read-only, idempotent]"),
5047            "got: {}",
5048            instructions
5049        );
5050    }
5051
5052    #[tokio::test]
5053    async fn test_auto_instructions_no_annotations_no_tags() {
5054        // Tools without annotations should have no tags at all
5055        let tool = ToolBuilder::new("fetch")
5056            .description("Fetch data")
5057            .handler(|_: AddInput| async move { Ok(CallToolResult::text("ok")) })
5058            .build();
5059
5060        let mut router = McpRouter::new().auto_instructions().tool(tool);
5061
5062        let resp = send_initialize(&mut router).await;
5063        let instructions = resp.instructions.unwrap();
5064
5065        // No bracket tags
5066        assert!(
5067            !instructions.contains('['),
5068            "should have no tags, got: {}",
5069            instructions
5070        );
5071        assert!(instructions.contains("- **fetch**: Fetch data"));
5072    }
5073
5074    /// Helper to send an Initialize request and return the result
5075    async fn send_initialize(router: &mut McpRouter) -> InitializeResult {
5076        let init_req = RouterRequest {
5077            id: RequestId::Number(0),
5078            inner: McpRequest::Initialize(InitializeParams {
5079                protocol_version: "2025-11-25".to_string(),
5080                capabilities: ClientCapabilities {
5081                    roots: None,
5082                    sampling: None,
5083                    elicitation: None,
5084                    tasks: None,
5085                    experimental: None,
5086                    extensions: None,
5087                },
5088                client_info: Implementation {
5089                    name: "test".to_string(),
5090                    version: "1.0".to_string(),
5091                    ..Default::default()
5092                },
5093                meta: None,
5094            }),
5095            extensions: Extensions::new(),
5096        };
5097        let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5098        match resp.inner {
5099            Ok(McpResponse::Initialize(result)) => result,
5100            other => panic!("Expected Initialize response, got {:?}", other),
5101        }
5102    }
5103
5104    #[tokio::test]
5105    async fn test_notify_tools_list_changed() {
5106        let (tx, mut rx) = crate::context::notification_channel(16);
5107
5108        let router = McpRouter::new()
5109            .server_info("test", "1.0")
5110            .with_notification_sender(tx);
5111
5112        assert!(router.notify_tools_list_changed());
5113
5114        let notification = rx.recv().await.unwrap();
5115        assert!(matches!(notification, ServerNotification::ToolsListChanged));
5116    }
5117
5118    #[tokio::test]
5119    async fn test_notify_prompts_list_changed() {
5120        let (tx, mut rx) = crate::context::notification_channel(16);
5121
5122        let router = McpRouter::new()
5123            .server_info("test", "1.0")
5124            .with_notification_sender(tx);
5125
5126        assert!(router.notify_prompts_list_changed());
5127
5128        let notification = rx.recv().await.unwrap();
5129        assert!(matches!(
5130            notification,
5131            ServerNotification::PromptsListChanged
5132        ));
5133    }
5134
5135    #[tokio::test]
5136    async fn test_notify_without_sender_returns_false() {
5137        let router = McpRouter::new().server_info("test", "1.0");
5138
5139        assert!(!router.notify_tools_list_changed());
5140        assert!(!router.notify_prompts_list_changed());
5141        assert!(!router.notify_resources_list_changed());
5142    }
5143
5144    #[tokio::test]
5145    async fn test_list_changed_capabilities_with_notification_sender() {
5146        let (tx, _rx) = crate::context::notification_channel(16);
5147        let tool = ToolBuilder::new("test")
5148            .description("test")
5149            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5150            .build();
5151
5152        let mut router = McpRouter::new()
5153            .server_info("test", "1.0")
5154            .tool(tool)
5155            .with_notification_sender(tx);
5156
5157        init_router(&mut router).await;
5158
5159        let caps = router.capabilities();
5160        let tools_cap = caps.tools.expect("tools capability should be present");
5161        assert!(
5162            tools_cap.list_changed,
5163            "tools.listChanged should be true when notification sender is configured"
5164        );
5165    }
5166
5167    #[tokio::test]
5168    async fn test_list_changed_capabilities_without_notification_sender() {
5169        let tool = ToolBuilder::new("test")
5170            .description("test")
5171            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5172            .build();
5173
5174        let mut router = McpRouter::new().server_info("test", "1.0").tool(tool);
5175
5176        init_router(&mut router).await;
5177
5178        let caps = router.capabilities();
5179        let tools_cap = caps.tools.expect("tools capability should be present");
5180        assert!(
5181            !tools_cap.list_changed,
5182            "tools.listChanged should be false without notification sender"
5183        );
5184    }
5185
5186    #[tokio::test]
5187    async fn test_set_logging_level_filters_messages() {
5188        let (tx, mut rx) = crate::context::notification_channel(16);
5189
5190        let mut router = McpRouter::new()
5191            .server_info("test", "1.0")
5192            .with_notification_sender(tx);
5193
5194        init_router(&mut router).await;
5195
5196        // Set logging level to Warning
5197        let set_level_req = RouterRequest {
5198            id: RequestId::Number(99),
5199            inner: McpRequest::SetLoggingLevel(SetLogLevelParams {
5200                level: LogLevel::Warning,
5201                meta: None,
5202            }),
5203            extensions: crate::context::Extensions::new(),
5204        };
5205        let resp = router
5206            .ready()
5207            .await
5208            .unwrap()
5209            .call(set_level_req)
5210            .await
5211            .unwrap();
5212        assert!(matches!(resp.inner, Ok(McpResponse::SetLoggingLevel(_))));
5213
5214        // Create a context from the router (simulating a handler)
5215        let ctx = router.create_context(RequestId::Number(100), None);
5216
5217        // Error (more severe than Warning) should pass through
5218        ctx.send_log(LoggingMessageParams::new(
5219            LogLevel::Error,
5220            serde_json::Value::Null,
5221        ));
5222        assert!(
5223            rx.try_recv().is_ok(),
5224            "Error should pass through Warning filter"
5225        );
5226
5227        // Info (less severe than Warning) should be filtered
5228        ctx.send_log(LoggingMessageParams::new(
5229            LogLevel::Info,
5230            serde_json::Value::Null,
5231        ));
5232        assert!(
5233            rx.try_recv().is_err(),
5234            "Info should be filtered at Warning level"
5235        );
5236    }
5237
5238    #[test]
5239    fn test_paginate_no_page_size() {
5240        let items = vec![1, 2, 3, 4, 5];
5241        let (page, cursor) = paginate(items.clone(), None, None).unwrap();
5242        assert_eq!(page, items);
5243        assert!(cursor.is_none());
5244    }
5245
5246    #[test]
5247    fn test_paginate_first_page() {
5248        let items = vec![1, 2, 3, 4, 5];
5249        let (page, cursor) = paginate(items, None, Some(2)).unwrap();
5250        assert_eq!(page, vec![1, 2]);
5251        assert!(cursor.is_some());
5252    }
5253
5254    #[test]
5255    fn test_paginate_middle_page() {
5256        let items = vec![1, 2, 3, 4, 5];
5257        let (page1, cursor1) = paginate(items.clone(), None, Some(2)).unwrap();
5258        assert_eq!(page1, vec![1, 2]);
5259
5260        let (page2, cursor2) = paginate(items, cursor1.as_deref(), Some(2)).unwrap();
5261        assert_eq!(page2, vec![3, 4]);
5262        assert!(cursor2.is_some());
5263    }
5264
5265    #[test]
5266    fn test_paginate_last_page() {
5267        let items = vec![1, 2, 3, 4, 5];
5268        // Skip to offset 4 (last item)
5269        let cursor = encode_cursor(4);
5270        let (page, next) = paginate(items, Some(&cursor), Some(2)).unwrap();
5271        assert_eq!(page, vec![5]);
5272        assert!(next.is_none());
5273    }
5274
5275    #[test]
5276    fn test_paginate_exact_boundary() {
5277        let items = vec![1, 2, 3, 4];
5278        let (page, cursor) = paginate(items, None, Some(4)).unwrap();
5279        assert_eq!(page, vec![1, 2, 3, 4]);
5280        assert!(cursor.is_none());
5281    }
5282
5283    #[test]
5284    fn test_paginate_invalid_cursor() {
5285        let items = vec![1, 2, 3];
5286        let result = paginate(items, Some("not-valid-base64!@#$"), Some(2));
5287        assert!(result.is_err());
5288    }
5289
5290    #[test]
5291    fn test_cursor_round_trip() {
5292        let offset = 42;
5293        let encoded = encode_cursor(offset);
5294        let decoded = decode_cursor(&encoded).unwrap();
5295        assert_eq!(decoded, offset);
5296    }
5297
5298    #[tokio::test]
5299    async fn test_list_tools_pagination() {
5300        let tool_a = ToolBuilder::new("alpha")
5301            .description("a")
5302            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5303            .build();
5304        let tool_b = ToolBuilder::new("beta")
5305            .description("b")
5306            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5307            .build();
5308        let tool_c = ToolBuilder::new("gamma")
5309            .description("c")
5310            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5311            .build();
5312
5313        let mut router = McpRouter::new()
5314            .server_info("test", "1.0")
5315            .page_size(2)
5316            .tool(tool_a)
5317            .tool(tool_b)
5318            .tool(tool_c);
5319
5320        init_router(&mut router).await;
5321
5322        // First page
5323        let req = RouterRequest {
5324            id: RequestId::Number(1),
5325            inner: McpRequest::ListTools(ListToolsParams {
5326                cursor: None,
5327                meta: None,
5328            }),
5329            extensions: Extensions::new(),
5330        };
5331        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5332        let (tools, next_cursor) = match resp.inner {
5333            Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5334            other => panic!("Expected ListTools, got {:?}", other),
5335        };
5336        assert_eq!(tools.len(), 2);
5337        assert_eq!(tools[0].name, "alpha");
5338        assert_eq!(tools[1].name, "beta");
5339        assert!(next_cursor.is_some());
5340
5341        // Second page
5342        let req = RouterRequest {
5343            id: RequestId::Number(2),
5344            inner: McpRequest::ListTools(ListToolsParams {
5345                cursor: next_cursor,
5346                meta: None,
5347            }),
5348            extensions: Extensions::new(),
5349        };
5350        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5351        let (tools, next_cursor) = match resp.inner {
5352            Ok(McpResponse::ListTools(result)) => (result.tools, result.next_cursor),
5353            other => panic!("Expected ListTools, got {:?}", other),
5354        };
5355        assert_eq!(tools.len(), 1);
5356        assert_eq!(tools[0].name, "gamma");
5357        assert!(next_cursor.is_none());
5358    }
5359
5360    #[tokio::test]
5361    async fn test_list_tools_no_pagination_by_default() {
5362        let tool_a = ToolBuilder::new("alpha")
5363            .description("a")
5364            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5365            .build();
5366        let tool_b = ToolBuilder::new("beta")
5367            .description("b")
5368            .handler(|_input: AddInput| async { Ok(CallToolResult::text("ok")) })
5369            .build();
5370
5371        let mut router = McpRouter::new()
5372            .server_info("test", "1.0")
5373            .tool(tool_a)
5374            .tool(tool_b);
5375
5376        init_router(&mut router).await;
5377
5378        let req = RouterRequest {
5379            id: RequestId::Number(1),
5380            inner: McpRequest::ListTools(ListToolsParams {
5381                cursor: None,
5382                meta: None,
5383            }),
5384            extensions: Extensions::new(),
5385        };
5386        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5387        match resp.inner {
5388            Ok(McpResponse::ListTools(result)) => {
5389                assert_eq!(result.tools.len(), 2);
5390                assert!(result.next_cursor.is_none());
5391            }
5392            other => panic!("Expected ListTools, got {:?}", other),
5393        }
5394    }
5395
5396    // =========================================================================
5397    // Dynamic Tool Registry Tests
5398    // =========================================================================
5399
5400    #[cfg(feature = "dynamic-tools")]
5401    mod dynamic_tools_tests {
5402        use super::*;
5403
5404        #[tokio::test]
5405        async fn test_dynamic_tools_register_and_list() {
5406            let (router, registry) = McpRouter::new()
5407                .server_info("test", "1.0")
5408                .with_dynamic_tools();
5409
5410            let tool = ToolBuilder::new("dynamic_echo")
5411                .description("Dynamic echo")
5412                .handler(|input: AddInput| async move {
5413                    Ok(CallToolResult::text(format!("{}", input.a)))
5414                })
5415                .build();
5416
5417            registry.register(tool);
5418
5419            let mut router = router;
5420            init_router(&mut router).await;
5421
5422            let req = RouterRequest {
5423                id: RequestId::Number(1),
5424                inner: McpRequest::ListTools(ListToolsParams::default()),
5425                extensions: Extensions::new(),
5426            };
5427
5428            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5429            match resp.inner {
5430                Ok(McpResponse::ListTools(result)) => {
5431                    assert_eq!(result.tools.len(), 1);
5432                    assert_eq!(result.tools[0].name, "dynamic_echo");
5433                }
5434                _ => panic!("Expected ListTools response"),
5435            }
5436        }
5437
5438        #[tokio::test]
5439        async fn test_dynamic_tools_unregister() {
5440            let (router, registry) = McpRouter::new()
5441                .server_info("test", "1.0")
5442                .with_dynamic_tools();
5443
5444            let tool = ToolBuilder::new("temp")
5445                .description("Temporary")
5446                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5447                .build();
5448
5449            registry.register(tool);
5450            assert!(registry.contains("temp"));
5451
5452            let removed = registry.unregister("temp");
5453            assert!(removed);
5454            assert!(!registry.contains("temp"));
5455
5456            // Unregistering again returns false
5457            assert!(!registry.unregister("temp"));
5458
5459            let mut router = router;
5460            init_router(&mut router).await;
5461
5462            let req = RouterRequest {
5463                id: RequestId::Number(1),
5464                inner: McpRequest::ListTools(ListToolsParams::default()),
5465                extensions: Extensions::new(),
5466            };
5467
5468            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5469            match resp.inner {
5470                Ok(McpResponse::ListTools(result)) => {
5471                    assert_eq!(result.tools.len(), 0);
5472                }
5473                _ => panic!("Expected ListTools response"),
5474            }
5475        }
5476
5477        #[tokio::test]
5478        async fn test_dynamic_tools_merged_with_static() {
5479            let static_tool = ToolBuilder::new("static_tool")
5480                .description("Static")
5481                .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5482                .build();
5483
5484            let (router, registry) = McpRouter::new()
5485                .server_info("test", "1.0")
5486                .tool(static_tool)
5487                .with_dynamic_tools();
5488
5489            let dynamic_tool = ToolBuilder::new("dynamic_tool")
5490                .description("Dynamic")
5491                .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5492                .build();
5493
5494            registry.register(dynamic_tool);
5495
5496            let mut router = router;
5497            init_router(&mut router).await;
5498
5499            let req = RouterRequest {
5500                id: RequestId::Number(1),
5501                inner: McpRequest::ListTools(ListToolsParams::default()),
5502                extensions: Extensions::new(),
5503            };
5504
5505            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5506            match resp.inner {
5507                Ok(McpResponse::ListTools(result)) => {
5508                    assert_eq!(result.tools.len(), 2);
5509                    let names: Vec<&str> = result.tools.iter().map(|t| t.name.as_str()).collect();
5510                    assert!(names.contains(&"static_tool"));
5511                    assert!(names.contains(&"dynamic_tool"));
5512                }
5513                _ => panic!("Expected ListTools response"),
5514            }
5515        }
5516
5517        #[tokio::test]
5518        async fn test_static_tools_shadow_dynamic() {
5519            let static_tool = ToolBuilder::new("shared")
5520                .description("Static version")
5521                .handler(|_: AddInput| async { Ok(CallToolResult::text("static")) })
5522                .build();
5523
5524            let (router, registry) = McpRouter::new()
5525                .server_info("test", "1.0")
5526                .tool(static_tool)
5527                .with_dynamic_tools();
5528
5529            let dynamic_tool = ToolBuilder::new("shared")
5530                .description("Dynamic version")
5531                .handler(|_: AddInput| async { Ok(CallToolResult::text("dynamic")) })
5532                .build();
5533
5534            registry.register(dynamic_tool);
5535
5536            let mut router = router;
5537            init_router(&mut router).await;
5538
5539            // List should only show the static version
5540            let req = RouterRequest {
5541                id: RequestId::Number(1),
5542                inner: McpRequest::ListTools(ListToolsParams::default()),
5543                extensions: Extensions::new(),
5544            };
5545
5546            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5547            match resp.inner {
5548                Ok(McpResponse::ListTools(result)) => {
5549                    assert_eq!(result.tools.len(), 1);
5550                    assert_eq!(result.tools[0].name, "shared");
5551                    assert_eq!(
5552                        result.tools[0].description.as_deref(),
5553                        Some("Static version")
5554                    );
5555                }
5556                _ => panic!("Expected ListTools response"),
5557            }
5558
5559            // Call should dispatch to the static tool
5560            let req = RouterRequest {
5561                id: RequestId::Number(2),
5562                inner: McpRequest::CallTool(CallToolParams {
5563                    name: "shared".to_string(),
5564                    arguments: serde_json::json!({"a": 1, "b": 2}),
5565                    meta: None,
5566                    task: None,
5567                }),
5568                extensions: Extensions::new(),
5569            };
5570
5571            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5572            match resp.inner {
5573                Ok(McpResponse::CallTool(result)) => {
5574                    assert!(!result.is_error);
5575                    match &result.content[0] {
5576                        Content::Text { text, .. } => assert_eq!(text, "static"),
5577                        _ => panic!("Expected text content"),
5578                    }
5579                }
5580                _ => panic!("Expected CallTool response"),
5581            }
5582        }
5583
5584        #[tokio::test]
5585        async fn test_dynamic_tools_call() {
5586            let (router, registry) = McpRouter::new()
5587                .server_info("test", "1.0")
5588                .with_dynamic_tools();
5589
5590            let tool = ToolBuilder::new("add")
5591                .description("Add two numbers")
5592                .handler(|input: AddInput| async move {
5593                    Ok(CallToolResult::text(format!("{}", input.a + input.b)))
5594                })
5595                .build();
5596
5597            registry.register(tool);
5598
5599            let mut router = router;
5600            init_router(&mut router).await;
5601
5602            let req = RouterRequest {
5603                id: RequestId::Number(1),
5604                inner: McpRequest::CallTool(CallToolParams {
5605                    name: "add".to_string(),
5606                    arguments: serde_json::json!({"a": 3, "b": 4}),
5607                    meta: None,
5608                    task: None,
5609                }),
5610                extensions: Extensions::new(),
5611            };
5612
5613            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5614            match resp.inner {
5615                Ok(McpResponse::CallTool(result)) => {
5616                    assert!(!result.is_error);
5617                    match &result.content[0] {
5618                        Content::Text { text, .. } => assert_eq!(text, "7"),
5619                        _ => panic!("Expected text content"),
5620                    }
5621                }
5622                _ => panic!("Expected CallTool response"),
5623            }
5624        }
5625
5626        #[tokio::test]
5627        async fn test_dynamic_tools_notification_on_register() {
5628            let (tx, mut rx) = crate::context::notification_channel(16);
5629            let (router, registry) = McpRouter::new()
5630                .server_info("test", "1.0")
5631                .with_dynamic_tools();
5632            let _router = router.with_notification_sender(tx);
5633
5634            let tool = ToolBuilder::new("notified")
5635                .description("Test")
5636                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5637                .build();
5638
5639            registry.register(tool);
5640
5641            let notification = rx.recv().await.unwrap();
5642            assert!(matches!(notification, ServerNotification::ToolsListChanged));
5643        }
5644
5645        #[tokio::test]
5646        async fn test_dynamic_tools_notification_on_unregister() {
5647            let (tx, mut rx) = crate::context::notification_channel(16);
5648            let (router, registry) = McpRouter::new()
5649                .server_info("test", "1.0")
5650                .with_dynamic_tools();
5651            let _router = router.with_notification_sender(tx);
5652
5653            let tool = ToolBuilder::new("notified")
5654                .description("Test")
5655                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5656                .build();
5657
5658            registry.register(tool);
5659            // Consume the register notification
5660            let _ = rx.recv().await.unwrap();
5661
5662            registry.unregister("notified");
5663            let notification = rx.recv().await.unwrap();
5664            assert!(matches!(notification, ServerNotification::ToolsListChanged));
5665        }
5666
5667        #[tokio::test]
5668        async fn test_dynamic_tools_no_notification_on_empty_unregister() {
5669            let (tx, mut rx) = crate::context::notification_channel(16);
5670            let (router, registry) = McpRouter::new()
5671                .server_info("test", "1.0")
5672                .with_dynamic_tools();
5673            let _router = router.with_notification_sender(tx);
5674
5675            // Unregister a tool that doesn't exist — should NOT send notification
5676            assert!(!registry.unregister("nonexistent"));
5677
5678            // Channel should be empty
5679            assert!(rx.try_recv().is_err());
5680        }
5681
5682        #[tokio::test]
5683        async fn test_dynamic_tools_filter_applies() {
5684            use crate::filter::CapabilityFilter;
5685
5686            let (router, registry) = McpRouter::new()
5687                .server_info("test", "1.0")
5688                .tool_filter(CapabilityFilter::new(|_, tool: &Tool| {
5689                    tool.name != "hidden"
5690                }))
5691                .with_dynamic_tools();
5692
5693            let visible = ToolBuilder::new("visible")
5694                .description("Visible")
5695                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5696                .build();
5697
5698            let hidden = ToolBuilder::new("hidden")
5699                .description("Hidden")
5700                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5701                .build();
5702
5703            registry.register(visible);
5704            registry.register(hidden);
5705
5706            let mut router = router;
5707            init_router(&mut router).await;
5708
5709            // List should only show visible tool
5710            let req = RouterRequest {
5711                id: RequestId::Number(1),
5712                inner: McpRequest::ListTools(ListToolsParams::default()),
5713                extensions: Extensions::new(),
5714            };
5715
5716            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5717            match resp.inner {
5718                Ok(McpResponse::ListTools(result)) => {
5719                    assert_eq!(result.tools.len(), 1);
5720                    assert_eq!(result.tools[0].name, "visible");
5721                }
5722                _ => panic!("Expected ListTools response"),
5723            }
5724
5725            // Call to hidden tool should be denied
5726            let req = RouterRequest {
5727                id: RequestId::Number(2),
5728                inner: McpRequest::CallTool(CallToolParams {
5729                    name: "hidden".to_string(),
5730                    arguments: serde_json::json!({"a": 1, "b": 2}),
5731                    meta: None,
5732                    task: None,
5733                }),
5734                extensions: Extensions::new(),
5735            };
5736
5737            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5738            match resp.inner {
5739                Err(e) => {
5740                    assert_eq!(e.code, -32601); // Method not found
5741                }
5742                _ => panic!("Expected JsonRpc error"),
5743            }
5744        }
5745
5746        #[tokio::test]
5747        async fn test_dynamic_tools_capabilities_advertised() {
5748            // No static tools, but dynamic tools enabled — should advertise tools capability
5749            let (mut router, _registry) = McpRouter::new()
5750                .server_info("test", "1.0")
5751                .with_dynamic_tools();
5752
5753            let init_req = RouterRequest {
5754                id: RequestId::Number(1),
5755                inner: McpRequest::Initialize(InitializeParams {
5756                    protocol_version: "2025-11-25".to_string(),
5757                    capabilities: ClientCapabilities::default(),
5758                    client_info: Implementation {
5759                        name: "test".to_string(),
5760                        version: "1.0".to_string(),
5761                        ..Default::default()
5762                    },
5763                    meta: None,
5764                }),
5765                extensions: Extensions::new(),
5766            };
5767
5768            let resp = router.ready().await.unwrap().call(init_req).await.unwrap();
5769            match resp.inner {
5770                Ok(McpResponse::Initialize(result)) => {
5771                    assert!(result.capabilities.tools.is_some());
5772                }
5773                _ => panic!("Expected Initialize response"),
5774            }
5775        }
5776
5777        #[tokio::test]
5778        async fn test_dynamic_tools_multi_session_notification() {
5779            let (tx1, mut rx1) = crate::context::notification_channel(16);
5780            let (tx2, mut rx2) = crate::context::notification_channel(16);
5781
5782            let (router, registry) = McpRouter::new()
5783                .server_info("test", "1.0")
5784                .with_dynamic_tools();
5785
5786            // Simulate two sessions by calling with_notification_sender on two clones
5787            let _session1 = router.clone().with_notification_sender(tx1);
5788            let _session2 = router.clone().with_notification_sender(tx2);
5789
5790            let tool = ToolBuilder::new("broadcast")
5791                .description("Test")
5792                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5793                .build();
5794
5795            registry.register(tool);
5796
5797            // Both sessions should receive the notification
5798            let n1 = rx1.recv().await.unwrap();
5799            let n2 = rx2.recv().await.unwrap();
5800            assert!(matches!(n1, ServerNotification::ToolsListChanged));
5801            assert!(matches!(n2, ServerNotification::ToolsListChanged));
5802        }
5803
5804        #[tokio::test]
5805        async fn test_dynamic_tools_call_not_found() {
5806            let (router, _registry) = McpRouter::new()
5807                .server_info("test", "1.0")
5808                .with_dynamic_tools();
5809
5810            let mut router = router;
5811            init_router(&mut router).await;
5812
5813            let req = RouterRequest {
5814                id: RequestId::Number(1),
5815                inner: McpRequest::CallTool(CallToolParams {
5816                    name: "nonexistent".to_string(),
5817                    arguments: serde_json::json!({}),
5818                    meta: None,
5819                    task: None,
5820                }),
5821                extensions: Extensions::new(),
5822            };
5823
5824            let resp = router.ready().await.unwrap().call(req).await.unwrap();
5825            match resp.inner {
5826                Err(e) => {
5827                    assert_eq!(e.code, -32601);
5828                }
5829                _ => panic!("Expected method not found error"),
5830            }
5831        }
5832
5833        #[tokio::test]
5834        async fn test_dynamic_tools_registry_list() {
5835            let (_, registry) = McpRouter::new()
5836                .server_info("test", "1.0")
5837                .with_dynamic_tools();
5838
5839            assert!(registry.list().is_empty());
5840
5841            let tool = ToolBuilder::new("tool_a")
5842                .description("A")
5843                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5844                .build();
5845            registry.register(tool);
5846
5847            let tool = ToolBuilder::new("tool_b")
5848                .description("B")
5849                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5850                .build();
5851            registry.register(tool);
5852
5853            let tools = registry.list();
5854            assert_eq!(tools.len(), 2);
5855            let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
5856            assert!(names.contains(&"tool_a"));
5857            assert!(names.contains(&"tool_b"));
5858        }
5859    } // mod dynamic_tools_tests
5860
5861    #[tokio::test]
5862    async fn test_tool_if_true_registers() {
5863        let tool = ToolBuilder::new("conditional")
5864            .description("Conditional tool")
5865            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5866            .build();
5867
5868        let mut router = McpRouter::new().tool_if(true, tool);
5869        init_router(&mut router).await;
5870
5871        let req = RouterRequest {
5872            id: RequestId::Number(1),
5873            inner: McpRequest::ListTools(ListToolsParams::default()),
5874            extensions: Extensions::new(),
5875        };
5876        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5877        match resp.inner {
5878            Ok(McpResponse::ListTools(result)) => {
5879                assert_eq!(result.tools.len(), 1);
5880                assert_eq!(result.tools[0].name, "conditional");
5881            }
5882            _ => panic!("Expected ListTools response"),
5883        }
5884    }
5885
5886    #[tokio::test]
5887    async fn test_tool_if_false_skips() {
5888        let tool = ToolBuilder::new("conditional")
5889            .description("Conditional tool")
5890            .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5891            .build();
5892
5893        let mut router = McpRouter::new().tool_if(false, tool);
5894        init_router(&mut router).await;
5895
5896        let req = RouterRequest {
5897            id: RequestId::Number(1),
5898            inner: McpRequest::ListTools(ListToolsParams::default()),
5899            extensions: Extensions::new(),
5900        };
5901        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5902        match resp.inner {
5903            Ok(McpResponse::ListTools(result)) => {
5904                assert_eq!(result.tools.len(), 0);
5905            }
5906            _ => panic!("Expected ListTools response"),
5907        }
5908    }
5909
5910    #[tokio::test]
5911    async fn test_tools_if_batch_conditional() {
5912        let tools = vec![
5913            ToolBuilder::new("a")
5914                .description("Tool A")
5915                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5916                .build(),
5917            ToolBuilder::new("b")
5918                .description("Tool B")
5919                .handler(|_: AddInput| async { Ok(CallToolResult::text("ok")) })
5920                .build(),
5921        ];
5922
5923        let mut router = McpRouter::new().tools_if(false, tools);
5924        init_router(&mut router).await;
5925
5926        let req = RouterRequest {
5927            id: RequestId::Number(1),
5928            inner: McpRequest::ListTools(ListToolsParams::default()),
5929            extensions: Extensions::new(),
5930        };
5931        let resp = router.ready().await.unwrap().call(req).await.unwrap();
5932        match resp.inner {
5933            Ok(McpResponse::ListTools(result)) => {
5934                assert_eq!(result.tools.len(), 0);
5935            }
5936            _ => panic!("Expected ListTools response"),
5937        }
5938    }
5939
5940    #[test]
5941    fn test_resource_if_true_registers() {
5942        let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
5943            .name("test")
5944            .text("hello");
5945
5946        let router = McpRouter::new().resource_if(true, resource);
5947        assert_eq!(router.inner.resources.len(), 1);
5948    }
5949
5950    #[test]
5951    fn test_resource_if_false_skips() {
5952        let resource = crate::resource::ResourceBuilder::new("file:///test.txt")
5953            .name("test")
5954            .text("hello");
5955
5956        let router = McpRouter::new().resource_if(false, resource);
5957        assert_eq!(router.inner.resources.len(), 0);
5958    }
5959
5960    #[test]
5961    fn test_prompt_if_true_registers() {
5962        let prompt = crate::prompt::PromptBuilder::new("greet")
5963            .description("Greeting")
5964            .user_message("Hello!");
5965
5966        let router = McpRouter::new().prompt_if(true, prompt);
5967        assert_eq!(router.inner.prompts.len(), 1);
5968    }
5969
5970    #[test]
5971    fn test_prompt_if_false_skips() {
5972        let prompt = crate::prompt::PromptBuilder::new("greet")
5973            .description("Greeting")
5974            .user_message("Hello!");
5975
5976        let router = McpRouter::new().prompt_if(false, prompt);
5977        assert_eq!(router.inner.prompts.len(), 0);
5978    }
5979}