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