Skip to main content

tower_mcp/
router.rs

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