Skip to main content

tower_mcp/
router.rs

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