diff --git a/crates/mcp-server/src/components.rs b/crates/mcp-server/src/components.rs index 7aa226ee..e362deac 100644 --- a/crates/mcp-server/src/components.rs +++ b/crates/mcp-server/src/components.rs @@ -8,11 +8,20 @@ use anyhow::Result; use futures::stream::{self, StreamExt}; use rmcp::model::{CallToolRequestParam, CallToolResult, Content, Tool}; use rmcp::{Peer, RoleServer}; -use serde_json::{json, Value}; +use serde_json::{json, Map, Value}; use tracing::{debug, error, info, instrument}; use wassette::schema::{canonicalize_output_schema, ensure_structured_result}; use wassette::{ComponentLoadOutcome, LifecycleManager, LoadResult}; +/// Extract tools filter from request arguments +fn extract_tools_filter(args: &Map) -> Option> { + args.get("tools").and_then(|v| v.as_array()).map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(|s| s.to_string())) + .collect::>() + }) +} + #[instrument(skip(lifecycle_manager))] pub(crate) async fn get_component_tools(lifecycle_manager: &LifecycleManager) -> Result> { debug!("Listing components"); @@ -52,13 +61,20 @@ pub(crate) async fn handle_load_component( .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required argument: 'path'"))?; + // Extract optional tools filter + let tools_filter = extract_tools_filter(&args); + debug!( path = %path, + tools_filter = ?tools_filter, operation = "load-component", "Component load operation started" ); - match lifecycle_manager.load_component(path).await { + match lifecycle_manager + .load_component_with_tools(path, tools_filter.as_deref()) + .await + { Ok(outcome) => { info!( path = %path, @@ -395,9 +411,15 @@ pub async fn handle_load_component_cli( .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("Missing required argument: 'path'"))?; - info!(path, "Loading component (CLI mode)"); + // Extract optional tools filter + let tools_filter = extract_tools_filter(&args); + + info!(path, tools_filter = ?tools_filter, "Loading component (CLI mode)"); - match lifecycle_manager.load_component(path).await { + match lifecycle_manager + .load_component_with_tools(path, tools_filter.as_deref()) + .await + { Ok(outcome) => { handle_tool_list_notification(None, &outcome.component_id, "load").await; create_load_component_success_result(&outcome) diff --git a/crates/mcp-server/src/tools.rs b/crates/mcp-server/src/tools.rs index d39a2c1a..6a594835 100644 --- a/crates/mcp-server/src/tools.rs +++ b/crates/mcp-server/src/tools.rs @@ -219,13 +219,21 @@ fn get_builtin_tools() -> Vec { Tool { name: Cow::Borrowed("load-component"), description: Some(Cow::Borrowed( - "Dynamically loads a new tool or component from either the filesystem or OCI registries.", + "Dynamically loads a new tool or component from either the filesystem or OCI registries. Optionally, you can specify which tools to load from the component.", )), input_schema: Arc::new( serde_json::from_value(json!({ "type": "object", "properties": { - "path": {"type": "string"} + "path": { + "type": "string", + "description": "Path to the component (file://, oci://, or https://)" + }, + "tools": { + "type": "array", + "items": {"type": "string"}, + "description": "Optional array of tool names to load from the component. If not specified, all tools will be loaded." + } }, "required": ["path"] })) diff --git a/crates/wassette/src/lib.rs b/crates/wassette/src/lib.rs index 8a9b24e3..083d1657 100644 --- a/crates/wassette/src/lib.rs +++ b/crates/wassette/src/lib.rs @@ -458,6 +458,7 @@ impl LifecycleManager { &self, component_id: &str, wasm_path: &Path, + tools_filter: Option<&[String]>, ) -> Result { let (component, wasm_bytes) = self .load_component_optimized(wasm_path, component_id) @@ -478,7 +479,7 @@ impl LifecycleManager { }; // Use package docs if available - let tool_metadata = if let Some(ref docs) = package_docs { + let mut tool_metadata = if let Some(ref docs) = package_docs { component_exports_to_tools_with_docs( &component_instance.component, self.runtime.as_ref(), @@ -489,6 +490,25 @@ impl LifecycleManager { component_exports_to_tools(&component_instance.component, self.runtime.as_ref(), true) }; + // Filter tools if a filter is provided + if let Some(filter) = tools_filter { + debug!( + component_id = %component_id, + filter = ?filter, + total_tools = tool_metadata.len(), + "Filtering tools for component" + ); + // Convert filter to HashSet for O(1) lookup performance + let filter_set: std::collections::HashSet<&str> = + filter.iter().map(|s| s.as_str()).collect(); + tool_metadata.retain(|tool| filter_set.contains(tool.normalized_name.as_str())); + info!( + component_id = %component_id, + filtered_count = tool_metadata.len(), + "Tools filtered successfully" + ); + } + let tool_names: Vec = tool_metadata .iter() .map(|tool| tool.normalized_name.clone()) @@ -526,13 +546,31 @@ impl LifecycleManager { /// component and whether it replaced an existing instance. #[instrument(skip(self))] pub async fn load_component(&self, uri: &str) -> Result { - debug!(uri, "Loading component"); + self.load_component_with_tools(uri, None).await + } + + /// Loads a new component from the given URI with optional tool filtering. + /// + /// This URI can be a file path, an OCI reference, or a URL. + /// If a component with the given id already exists, it will be updated with the new component. + /// The `tools_filter` parameter allows you to specify which tools to load from the component. + /// If `None`, all tools will be loaded. + /// + /// Returns rich [`ComponentLoadOutcome`] information describing the loaded + /// component and whether it replaced an existing instance. + #[instrument(skip(self))] + pub async fn load_component_with_tools( + &self, + uri: &str, + tools_filter: Option<&[String]>, + ) -> Result { + debug!(uri, tools_filter = ?tools_filter, "Loading component"); let (component_id, resource) = self.resolve_component_resource(uri).await?; let staged_path = self .stage_component_artifact(&component_id, resource) .await?; let outcome = self - .compile_and_register_component(&component_id, &staged_path) + .compile_and_register_component(&component_id, &staged_path, tools_filter) .await .with_context(|| { format!( @@ -829,7 +867,7 @@ impl LifecycleManager { bail!("Component not found: {}", component_id); } - self.compile_and_register_component(component_id, &entry_path) + self.compile_and_register_component(component_id, &entry_path, None) .await .with_context(|| { format!( @@ -1251,7 +1289,7 @@ impl LifecycleManager { } let start_time = Instant::now(); - self.compile_and_register_component(&component_id, &entry_path) + self.compile_and_register_component(&component_id, &entry_path, None) .await .with_context(|| { format!( diff --git a/docs/reference/built-in-tools.md b/docs/reference/built-in-tools.md index 0cc69872..038c889a 100644 --- a/docs/reference/built-in-tools.md +++ b/docs/reference/built-in-tools.md @@ -4,7 +4,7 @@ Wassette comes with several built-in tools for managing components and their per | Tool | Description | |------|-------------| -| `load-component` | Dynamically loads a new tool or component from either the filesystem or OCI registries | +| `load-component` | Dynamically loads a new tool or component from either the filesystem or OCI registries. Optionally, you can specify which tools to load from the component | | `unload-component` | Unloads a tool or component | | `list-components` | Lists all currently loaded components or tools | | `search-components` | Lists all known components that can be fetched and loaded from the component registry | @@ -23,6 +23,7 @@ Wassette comes with several built-in tools for managing components and their per ## load-component **Parameters:** - `path` (string, required): Path to the component from either filesystem or OCI registries (e.g., `oci://ghcr.io/microsoft/time-server-js:latest` or `/path/to/component.wasm`) +- `tools` (array of strings, optional): Array of tool names to load from the component. If not specified, all tools will be loaded. This allows you to selectively load only specific tools from a component. **Returns:** ```json @@ -35,6 +36,23 @@ Wassette comes with several built-in tools for managing components and their per When an existing component is replaced, the `status` value becomes `component reloaded successfully`. +**Examples:** + +Load all tools from a component: +```json +{ + "path": "oci://ghcr.io/microsoft/fetch-rs:latest" +} +``` + +Load only specific tools from a component: +```json +{ + "path": "oci://ghcr.io/microsoft/fetch-rs:latest", + "tools": ["fetch"] +} +``` + ## unload-component **Parameters:** - `id` (string, required): Unique identifier of the component to unload diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 0beb1081..d0961b8e 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -124,8 +124,18 @@ wassette component load file:///path/to/component.wasm wassette component load file://./my-component.wasm ``` +**Load only specific tools from a component:** +```bash +# Load only the 'fetch' tool from the fetch-rs component +wassette component load oci://ghcr.io/microsoft/fetch-rs:latest --tools fetch + +# Load multiple specific tools +wassette component load file:///path/to/component.wasm --tools tool1,tool2,tool3 +``` + **Options:** - `--component-dir `: Component storage directory +- `--tools `: Comma-separated list of tool names to load from the component. If not specified, all tools will be loaded. ### `wassette component unload` diff --git a/src/commands.rs b/src/commands.rs index 65240889..d7372212 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -177,6 +177,9 @@ pub enum ComponentCommands { /// Directory where components are stored. Defaults to $XDG_DATA_HOME/wassette/components #[arg(long)] component_dir: Option, + /// Optional comma-separated list of tool names to load from the component. If not specified, all tools will be loaded. + #[arg(long, value_delimiter = ',')] + tools: Option>, }, /// Unload a WebAssembly component. Unload { diff --git a/src/main.rs b/src/main.rs index 5601161b..d61e47cd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -274,11 +274,15 @@ async fn main() -> Result<()> { ComponentCommands::Load { path, component_dir, + tools, } => { let component_dir = component_dir.clone().or_else(|| cli.component_dir.clone()); let lifecycle_manager = create_lifecycle_manager(component_dir).await?; let mut args = Map::new(); args.insert("path".to_string(), json!(path)); + if let Some(tools_list) = tools { + args.insert("tools".to_string(), json!(tools_list)); + } handle_tool_cli_command( &lifecycle_manager, "load-component", diff --git a/tests/selective_tool_loading_test.rs b/tests/selective_tool_loading_test.rs new file mode 100644 index 00000000..22608311 --- /dev/null +++ b/tests/selective_tool_loading_test.rs @@ -0,0 +1,163 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Integration tests for selective tool loading + +use std::path::PathBuf; + +use anyhow::{Context, Result}; +use wassette::LifecycleManager; + +/// Helper function to get the path to the fetch-rs test component +fn get_fetch_component_path() -> Result { + let project_root = std::env::current_dir().context("Failed to get current directory")?; + Ok(project_root + .join("examples") + .join("fetch-rs") + .join("target") + .join("wasm32-wasip2") + .join("release") + .join("fetch_rs.wasm")) +} + +#[tokio::test] +async fn test_selective_tool_loading() -> Result<()> { + // Create lifecycle manager + let tempdir = tempfile::tempdir()?; + let lifecycle_manager = LifecycleManager::new(&tempdir).await?; + + let fetch_component_path = get_fetch_component_path()?; + + // Load component with only specific tools + let tools_to_load = vec!["fetch".to_string()]; + let outcome = lifecycle_manager + .load_component_with_tools( + &format!("file://{}", fetch_component_path.display()), + Some(&tools_to_load), + ) + .await + .context("Failed to load component with selective tools")?; + + // Verify only the specified tool was loaded + assert_eq!( + outcome.tool_names.len(), + 1, + "Expected only 1 tool to be loaded" + ); + assert_eq!( + outcome.tool_names[0], "fetch", + "Expected 'fetch' tool to be loaded" + ); + + // Verify we can list the component and it shows the filtered tools + let components = lifecycle_manager.list_components().await; + assert!( + components.contains(&outcome.component_id), + "Component should be in the list" + ); + + // Get the component schema and verify it only has the filtered tools + let schema = lifecycle_manager + .get_component_schema(&outcome.component_id) + .await + .context("Failed to get component schema")?; + + let tools = schema + .get("tools") + .and_then(|v| v.as_array()) + .context("Schema should have tools array")?; + + assert_eq!( + tools.len(), + 1, + "Schema should only contain 1 tool after filtering" + ); + + // Verify the tool in the schema is 'fetch' + let tool_name = tools[0] + .get("name") + .and_then(|v| v.as_str()) + .context("Tool should have name")?; + assert_eq!(tool_name, "fetch", "Tool name should be 'fetch'"); + + Ok(()) +} + +#[tokio::test] +async fn test_load_all_tools_when_no_filter() -> Result<()> { + // Create lifecycle manager + let tempdir = tempfile::tempdir()?; + let lifecycle_manager = LifecycleManager::new(&tempdir).await?; + + let fetch_component_path = get_fetch_component_path()?; + + // Load component without filter (should load all tools) + let outcome = lifecycle_manager + .load_component(&format!("file://{}", fetch_component_path.display())) + .await + .context("Failed to load component")?; + + // The fetch-rs component should have at least one tool + assert!( + !outcome.tool_names.is_empty(), + "Component should have at least one tool" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_empty_tools_filter() -> Result<()> { + // Create lifecycle manager + let tempdir = tempfile::tempdir()?; + let lifecycle_manager = LifecycleManager::new(&tempdir).await?; + + let fetch_component_path = get_fetch_component_path()?; + + // Load component with empty filter (should load no tools) + let tools_to_load: Vec = vec![]; + let outcome = lifecycle_manager + .load_component_with_tools( + &format!("file://{}", fetch_component_path.display()), + Some(&tools_to_load), + ) + .await + .context("Failed to load component with empty tools filter")?; + + // Verify no tools were loaded + assert_eq!( + outcome.tool_names.len(), + 0, + "Expected 0 tools to be loaded with empty filter" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_nonexistent_tool_filter() -> Result<()> { + // Create lifecycle manager + let tempdir = tempfile::tempdir()?; + let lifecycle_manager = LifecycleManager::new(&tempdir).await?; + + let fetch_component_path = get_fetch_component_path()?; + + // Load component with filter for non-existent tool + let tools_to_load = vec!["nonexistent_tool".to_string()]; + let outcome = lifecycle_manager + .load_component_with_tools( + &format!("file://{}", fetch_component_path.display()), + Some(&tools_to_load), + ) + .await + .context("Failed to load component")?; + + // Verify no tools were loaded (filter matched nothing) + assert_eq!( + outcome.tool_names.len(), + 0, + "Expected 0 tools when filter doesn't match any tools" + ); + + Ok(()) +}