Add selective tool loading to component load

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<String, Value>) -> Option<Vec<String>> {
+    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::<Vec<String>>()
+    })
+}
+
 #[instrument(skip(lifecycle_manager))]
 pub(crate) async fn get_component_tools(lifecycle_manager: &LifecycleManager) -> Result<Vec<Tool>> {
     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)

crates/mcp-server/src/tools.rs

@@ -219,13 +219,21 @@ fn get_builtin_tools() -> Vec<Tool> {
         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"]
                 }))

crates/wassette/src/lib.rs

@@ -458,6 +458,7 @@ impl LifecycleManager {
         &self,
         component_id: &str,
         wasm_path: &Path,
+        tools_filter: Option<&[String]>,
     ) -> Result<ComponentLoadOutcome> {
         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<String> = 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<ComponentLoadOutcome> {
-        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<ComponentLoadOutcome> {
+        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!(

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

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 <PATH>`: Component storage directory
+- `--tools <TOOLS>`: Comma-separated list of tool names to load from the component. If not specified, all tools will be loaded.
 
 ### `wassette component unload`
 

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<PathBuf>,
+        /// 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<Vec<String>>,
     },
     /// Unload a WebAssembly component.
     Unload {

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",