From 8a60fa2d37793e27d797b4438feec81e8ed8330a Mon Sep 17 00:00:00 2001 From: Henry Park Date: Thu, 12 Mar 2026 16:30:38 -0700 Subject: [PATCH] fix: add tool_info schema discovery for WASM tools (#1086) * fix: add tool_info schema discovery for WASM tools * refactor: simplify WASM schema and hint state * refactor: store tool_info registry reference as Weak --- src/app.rs | 1 + src/tools/builtin/mod.rs | 2 + src/tools/builtin/tool_info.rs | 183 +++++++++++ src/tools/registry.rs | 12 + src/tools/tool.rs | 11 + src/tools/wasm/error.rs | 90 +----- src/tools/wasm/limits.rs | 8 - src/tools/wasm/mod.rs | 2 +- src/tools/wasm/runtime.rs | 67 ++-- src/tools/wasm/wrapper.rs | 300 ++++++++++++++---- tests/e2e_builtin_tool_coverage.rs | 86 +++++ .../llm_traces/tools/tool_info_discovery.json | 50 +++ .../web-search-tool.capabilities.json | 35 ++ 13 files changed, 658 insertions(+), 189 deletions(-) create mode 100644 src/tools/builtin/tool_info.rs create mode 100644 tests/fixtures/llm_traces/tools/tool_info_discovery.json diff --git a/src/app.rs b/src/app.rs index 59890bfc..da77d3f3 100644 --- a/src/app.rs +++ b/src/app.rs @@ -290,6 +290,7 @@ impl AppBuilder { Arc::new(ToolRegistry::new()) }; tools.register_builtin_tools(); + tools.register_tool_info(); if let Some(ref ss) = self.secrets_store { tools.register_secrets_tools(Arc::clone(ss)); diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index b52502c9..8ba8e57b 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -15,6 +15,7 @@ pub mod secrets_tools; pub(crate) mod shell; pub mod skill_tools; mod time; +mod tool_info; pub use echo::EchoTool; pub use extension_tools::{ @@ -39,6 +40,7 @@ pub use secrets_tools::{SecretDeleteTool, SecretListTool}; pub use shell::ShellTool; pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool}; pub use time::TimeTool; +pub use tool_info::ToolInfoTool; mod html_converter; pub mod image_analyze; pub mod image_edit; diff --git a/src/tools/builtin/tool_info.rs b/src/tools/builtin/tool_info.rs new file mode 100644 index 00000000..cd94384d --- /dev/null +++ b/src/tools/builtin/tool_info.rs @@ -0,0 +1,183 @@ +//! On-demand tool discovery (like CLI `--help`). +//! +//! Two levels of detail: +//! - Default: name, description, parameter names (compact ~150 bytes) +//! - `include_schema: true`: adds the full typed JSON Schema +//! +//! Keeps the tools array compact (WASM tools use permissive schemas) +//! while allowing precise discovery when needed. + +use std::sync::Weak; + +use async_trait::async_trait; + +use crate::context::JobContext; +use crate::tools::registry::ToolRegistry; +use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str}; + +pub struct ToolInfoTool { + registry: Weak, +} + +impl ToolInfoTool { + pub fn new(registry: Weak) -> Self { + Self { registry } + } +} + +#[async_trait] +impl Tool for ToolInfoTool { + fn name(&self) -> &str { + "tool_info" + } + + fn description(&self) -> &str { + "Get info about any tool: description and parameter names. \ + Set include_schema to true for the full typed parameter schema." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the tool to get info about" + }, + "include_schema": { + "type": "boolean", + "description": "If true, include the full typed JSON Schema for parameters (larger response). Default: false.", + "default": false + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + let name = require_str(¶ms, "name")?; + let include_schema = params + .get("include_schema") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let registry = self.registry.upgrade().ok_or_else(|| { + ToolError::ExecutionFailed( + "tool registry is no longer available for tool_info".to_string(), + ) + })?; + + let tool = registry.get(name).await.ok_or_else(|| { + ToolError::InvalidParameters(format!("No tool named '{name}' is registered")) + })?; + + let schema = tool.discovery_schema(); + + // Extract just param names from the schema's "properties" keys + let param_names: Vec<&str> = schema + .get("properties") + .and_then(|p| p.as_object()) + .map(|props| props.keys().map(|k| k.as_str()).collect()) + .unwrap_or_default(); + + let mut info = serde_json::json!({ + "name": tool.name(), + "description": tool.description(), + "parameters": param_names, + }); + + if include_schema { + info["schema"] = schema; + } + + Ok(ToolOutput::success(info, start.elapsed())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::builtin::EchoTool; + use std::sync::Arc; + + #[tokio::test] + async fn test_tool_info_default_returns_param_names() { + let registry = Arc::new(ToolRegistry::new()); + registry.register(Arc::new(EchoTool)).await; + + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "echo"}), &ctx) + .await + .unwrap(); + + let info = &result.result; + assert_eq!(info["name"], "echo"); + assert!(!info["description"].as_str().unwrap().is_empty()); + // Default: parameters is an array of names, not the full schema + assert!(info["parameters"].is_array()); + assert!( + info["parameters"] + .as_array() + .unwrap() + .iter() + .any(|v| v.as_str() == Some("message")), + "echo tool should have 'message' parameter: {:?}", + info["parameters"] + ); + // No schema field by default + assert!(info.get("schema").is_none()); + } + + #[tokio::test] + async fn test_tool_info_with_schema() { + let registry = Arc::new(ToolRegistry::new()); + registry.register(Arc::new(EchoTool)).await; + + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute( + serde_json::json!({"name": "echo", "include_schema": true}), + &ctx, + ) + .await + .unwrap(); + + let info = &result.result; + assert_eq!(info["name"], "echo"); + // With include_schema: true, schema field should be present + assert!(info["schema"].is_object()); + assert!(info["schema"]["properties"].is_object()); + } + + #[tokio::test] + async fn test_tool_info_unknown_tool() { + let registry = Arc::new(ToolRegistry::new()); + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "nonexistent"}), &ctx) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_tool_info_registry_dropped() { + let registry = Arc::new(ToolRegistry::new()); + let tool = ToolInfoTool::new(Arc::downgrade(®istry)); + drop(registry); + + let ctx = JobContext::default(); + let result = tool + .execute(serde_json::json!({"name": "echo"}), &ctx) + .await; + assert!(matches!(result, Err(ToolError::ExecutionFailed(_)))); + } +} diff --git a/src/tools/registry.rs b/src/tools/registry.rs index bee8cf27..754869c8 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -75,6 +75,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "image_generate", "image_edit", "image_analyze", + "tool_info", ]; /// Registry of available tools. @@ -245,6 +246,17 @@ impl ToolRegistry { tracing::debug!("Registered {} built-in tools", self.count()); } + /// Register the `tool_info` discovery tool. + /// + /// Requires `Arc` so the tool can query the registry for other tools' + /// schemas at runtime. Call after `register_builtin_tools()`. + pub fn register_tool_info(self: &Arc) { + use crate::tools::builtin::ToolInfoTool; + let tool = ToolInfoTool::new(Arc::downgrade(self)); + self.register_sync(Arc::new(tool)); + tracing::debug!("Registered tool_info discovery tool"); + } + /// Register only orchestrator-domain tools (safe for the main process). /// /// This registers tools that don't touch the filesystem or run shell commands: diff --git a/src/tools/tool.rs b/src/tools/tool.rs index b5879e5a..4a0fda8d 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -336,6 +336,17 @@ pub trait Tool: Send + Sync { None } + /// Full parameter schema for discovery and coercion purposes. + /// + /// Unlike `parameters_schema()` (which may be permissive to keep the tools + /// array compact), this returns the complete typed schema. Used by the + /// `tool_info` built-in and by WASM parameter coercion. + /// + /// Default: delegates to `parameters_schema()`. + fn discovery_schema(&self) -> serde_json::Value { + self.parameters_schema() + } + /// Get the tool schema for LLM function calling. fn schema(&self) -> ToolSchema { ToolSchema { diff --git a/src/tools/wasm/error.rs b/src/tools/wasm/error.rs index 8bbb8202..a0900775 100644 --- a/src/tools/wasm/error.rs +++ b/src/tools/wasm/error.rs @@ -1,7 +1,5 @@ //! WASM sandbox error types. -use std::fmt; - use thiserror::Error; /// Errors that can occur during WASM tool execution. @@ -68,13 +66,13 @@ pub enum WasmError { Timeout(std::time::Duration), /// Component returned an error response. - /// When `hint` is non-empty it carries the tool's description and parameter - /// schema so the LLM can retry with correct arguments. + /// When `hint` is non-empty it points the LLM to `tool_info` so it can + /// fetch the tool's full parameter schema on demand. #[error("Tool error: {message}{}", if hint.is_empty() { String::new() } else { format!("\n\nTool usage hint:\n{hint}") })] ToolReturnedError { /// The error message from the WASM tool. message: String, - /// Optional description + schema hint (empty when unavailable). + /// Optional retry hint (empty when unavailable). hint: String, }, @@ -99,73 +97,9 @@ impl From for crate::tools::ToolError { } } -/// Details about a trap that occurred during execution. -#[derive(Debug, Clone)] -pub struct TrapInfo { - /// Human-readable trap message. - pub message: String, - /// Trap code if available. - pub code: Option, -} - -impl fmt::Display for TrapInfo { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match &self.code { - Some(code) => write!(f, "{}: {}", code, self.message), - None => write!(f, "{}", self.message), - } - } -} - -/// Known trap codes from Wasmtime. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TrapCode { - /// Out of bounds memory access. - MemoryOutOfBounds, - /// Out of bounds table access. - TableOutOfBounds, - /// Indirect call type mismatch. - IndirectCallToNull, - /// Signature mismatch on indirect call. - BadSignature, - /// Integer overflow. - IntegerOverflow, - /// Integer division by zero. - IntegerDivisionByZero, - /// Invalid conversion to integer. - BadConversionToInteger, - /// Unreachable instruction executed. - UnreachableCodeReached, - /// Call stack exhausted. - StackOverflow, - /// Out of fuel. - OutOfFuel, - /// Unknown trap code. - Unknown, -} - -impl fmt::Display for TrapCode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let s = match self { - TrapCode::MemoryOutOfBounds => "memory out of bounds", - TrapCode::TableOutOfBounds => "table out of bounds", - TrapCode::IndirectCallToNull => "indirect call to null", - TrapCode::BadSignature => "bad signature", - TrapCode::IntegerOverflow => "integer overflow", - TrapCode::IntegerDivisionByZero => "integer division by zero", - TrapCode::BadConversionToInteger => "bad conversion to integer", - TrapCode::UnreachableCodeReached => "unreachable code reached", - TrapCode::StackOverflow => "stack overflow", - TrapCode::OutOfFuel => "out of fuel", - TrapCode::Unknown => "unknown trap", - }; - write!(f, "{}", s) - } -} - #[cfg(test)] mod tests { - use crate::tools::wasm::error::{TrapCode, TrapInfo, WasmError}; + use crate::tools::wasm::error::WasmError; #[test] fn test_error_display() { @@ -180,17 +114,6 @@ mod tests { assert!(err.to_string().contains("10000000")); } - #[test] - fn test_trap_info_display() { - let info = TrapInfo { - message: "access at offset 0x1000".to_string(), - code: Some(TrapCode::MemoryOutOfBounds), - }; - let s = info.to_string(); - assert!(s.contains("memory out of bounds")); - assert!(s.contains("access at offset")); - } - #[test] fn test_conversion_to_tool_error() { let wasm_err = WasmError::Trapped("test trap".to_string()); @@ -218,12 +141,11 @@ mod tests { fn test_tool_returned_error_with_hint() { let err = WasmError::ToolReturnedError { message: "unknown action: foobar".to_string(), - hint: "Description: Gmail tool\nParameters schema: {\"type\":\"object\"}".to_string(), + hint: "Tip: call tool_info(name: \"gmail\", include_schema: true) for the full parameter schema.".to_string(), }; let display = err.to_string(); assert!(display.contains("unknown action: foobar")); assert!(display.contains("Tool usage hint")); - assert!(display.contains("Gmail tool")); - assert!(display.contains("Parameters schema")); + assert!(display.contains("tool_info")); } } diff --git a/src/tools/wasm/limits.rs b/src/tools/wasm/limits.rs index 237247e9..d537a583 100644 --- a/src/tools/wasm/limits.rs +++ b/src/tools/wasm/limits.rs @@ -67,14 +67,8 @@ pub struct WasmResourceLimiter { memory_used: u64, /// Maximum tables allowed. max_tables: u32, - /// Current table count. - #[allow(dead_code)] // Reserved for table limit enforcement - tables_created: u32, /// Maximum instances allowed. max_instances: u32, - /// Current instance count. - #[allow(dead_code)] // Reserved for instance limit enforcement - instances_created: u32, } impl WasmResourceLimiter { @@ -87,9 +81,7 @@ impl WasmResourceLimiter { memory_limit, memory_used: 0, max_tables: 10, - tables_created: 0, max_instances: 10, // Component model needs multiple instances for WASI - instances_created: 0, } } diff --git a/src/tools/wasm/mod.rs b/src/tools/wasm/mod.rs index 647d2837..1998e801 100644 --- a/src/tools/wasm/mod.rs +++ b/src/tools/wasm/mod.rs @@ -96,7 +96,7 @@ pub(crate) mod storage; mod wrapper; // Core types -pub use error::{TrapCode, TrapInfo, WasmError}; +pub use error::WasmError; pub use host::{HostState, LogEntry, LogLevel}; pub use limits::{ DEFAULT_FUEL_LIMIT, DEFAULT_MEMORY_LIMIT, DEFAULT_TIMEOUT, FuelConfig, ResourceLimits, diff --git a/src/tools/wasm/runtime.rs b/src/tools/wasm/runtime.rs index 7645af3e..02c56f61 100644 --- a/src/tools/wasm/runtime.rs +++ b/src/tools/wasm/runtime.rs @@ -123,7 +123,9 @@ pub struct PreparedModule { pub name: String, /// Tool description (cached from component). pub description: String, - /// Parameter schema JSON (cached from component). + /// Full parameter schema JSON extracted from the component. + /// Used for discovery and coercion, not necessarily for the compact + /// schema advertised in the main tools array. pub schema: serde_json::Value, /// Pre-compiled component (cheaply cloneable via internal Arc). component: wasmtime::component::Component, @@ -265,11 +267,29 @@ impl WasmToolRuntime { let component = wasmtime::component::Component::new(&engine, &wasm_bytes) .map_err(|e| WasmError::CompilationFailed(e.to_string()))?; - // We need to instantiate briefly to extract metadata. - // In a full implementation, we'd use WIT bindgen to get typed access. - // For now, we extract what we can from the component. - let description = extract_tool_description(&engine, &component)?; - let schema = extract_tool_schema(&engine, &component)?; + // Briefly instantiate to extract metadata (description + schema) + // from the tool's exports, analogous to MCP's list_tools(). + let effective_limits = limits.clone().unwrap_or(default_limits.clone()); + let (description, schema) = crate::tools::wasm::wrapper::extract_wasm_metadata( + &engine, + &component, + &effective_limits, + ) + .unwrap_or_else(|e| { + tracing::warn!( + name = %name, + error = %e, + "WASM metadata extraction failed, using fallbacks" + ); + ( + "WASM sandboxed tool".to_string(), + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }), + ) + }); Ok::<_, WasmError>(PreparedModule { name: name.clone(), @@ -321,41 +341,6 @@ impl WasmToolRuntime { } } -/// Extract tool description from a compiled component. -/// -/// Returns a generic fallback. Callers should prefer loading the description -/// from the sidecar `*.capabilities.json` file and overriding via -/// `WasmToolWrapper::with_description()` or the `WasmToolRegistration::description` field. -fn extract_tool_description( - _engine: &Engine, - _component: &wasmtime::component::Component, -) -> Result { - // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). - // Real descriptions come from the capabilities.json sidecar file, which is - // loaded by the WasmToolLoader and passed as an override at registration time. - Ok("WASM sandboxed tool".to_string()) -} - -/// Extract tool parameter schema from a compiled component. -/// -/// Returns a permissive fallback that accepts any JSON object. Callers should -/// prefer loading the schema from the sidecar `*.capabilities.json` file and -/// overriding via `WasmToolWrapper::with_schema()` or the -/// `WasmToolRegistration::schema` field. -fn extract_tool_schema( - _engine: &Engine, - _component: &wasmtime::component::Component, -) -> Result { - // WIT bindgen extraction is not yet implemented (see TODO #4 in CLAUDE.md). - // Real schemas come from the capabilities.json sidecar file, which is - // loaded by the WasmToolLoader and passed as an override at registration time. - Ok(serde_json::json!({ - "type": "object", - "properties": {}, - "additionalProperties": true - })) -} - impl std::fmt::Debug for WasmToolRuntime { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("WasmToolRuntime") diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index c0294c51..52805afa 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -465,8 +465,8 @@ pub struct WasmToolWrapper { capabilities: Capabilities, /// Cached description (from PreparedModule or override). description: String, - /// Cached schema (from PreparedModule or override). - schema: serde_json::Value, + /// Compact and discovery schemas for this tool. + schemas: WasmToolSchemas, /// Injected credentials for HTTP requests (e.g., OAuth tokens). /// Keys are placeholder names like "GOOGLE_ACCESS_TOKEN". credentials: HashMap, @@ -477,6 +477,79 @@ pub struct WasmToolWrapper { oauth_refresh: Option, } +#[derive(Debug, Clone)] +struct WasmToolSchemas { + /// Compact schema advertised in the main tools array. + /// + /// This stays permissive by default to avoid serializing full exported + /// WASM schemas on every LLM call. Sidecars can override it explicitly. + advertised: serde_json::Value, + /// Full schema available for discovery and coercion. + /// + /// Seeded from the WASM `schema()` export at registration time, unless a + /// sidecar explicitly overrides it. + discovery: serde_json::Value, +} + +impl WasmToolSchemas { + fn permissive_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }) + } + + fn is_permissive_schema(schema: &serde_json::Value) -> bool { + schema + .get("properties") + .and_then(|p| p.as_object()) + .is_none_or(|p| p.is_empty()) + } + + fn new(discovery: serde_json::Value) -> Self { + Self { + advertised: Self::permissive_schema(), + discovery, + } + } + + fn with_override(&self, schema: serde_json::Value) -> Self { + Self { + advertised: schema.clone(), + discovery: schema, + } + } + + fn is_advertised_permissive(&self) -> bool { + Self::is_permissive_schema(&self.advertised) + } + + fn advertised(&self) -> serde_json::Value { + self.advertised.clone() + } + + fn discovery(&self) -> serde_json::Value { + self.discovery.clone() + } + + fn effective_for_coercion( + &self, + tool_iface: &wit_tool::Guest, + store: &mut Store, + ) -> serde_json::Value { + if !Self::is_permissive_schema(&self.discovery) { + return self.discovery.clone(); + } + + tool_iface + .call_schema(store) + .ok() + .and_then(|schema_str| serde_json::from_str::(&schema_str).ok()) + .unwrap_or_else(|| self.discovery.clone()) + } +} + impl WasmToolWrapper { /// Create a new WASM tool wrapper. pub fn new( @@ -484,30 +557,54 @@ impl WasmToolWrapper { prepared: Arc, capabilities: Capabilities, ) -> Self { - Self { + let mut wrapper = Self { description: prepared.description.clone(), - schema: prepared.schema.clone(), + schemas: WasmToolSchemas::new(prepared.schema.clone()), runtime, prepared, capabilities, credentials: HashMap::new(), secrets_store: None, oauth_refresh: None, - } + }; + wrapper.append_schema_hint_if_permissive(); + wrapper } /// Override the tool description. pub fn with_description(mut self, description: impl Into) -> Self { self.description = description.into(); + self.append_schema_hint_if_permissive(); self } /// Override the parameter schema. pub fn with_schema(mut self, schema: serde_json::Value) -> Self { - self.schema = schema; + self.schemas = self.schemas.with_override(schema); + self.strip_schema_hint(); + self.append_schema_hint_if_permissive(); self } + /// Append a tool_info hint to the description when the schema is permissive + /// (no typed properties), so the LLM knows to call tool_info for the full schema. + fn append_schema_hint_if_permissive(&mut self) { + if self.schemas.is_advertised_permissive() && !self.description.contains("tool_info") { + self.description + .push_str(" (call tool_info for parameter schema)"); + } + } + + /// Remove the tool_info hint from the description (e.g. after with_schema adds real types). + fn strip_schema_hint(&mut self) { + if let Some(pos) = self + .description + .find(" (call tool_info for parameter schema)") + { + self.description.truncate(pos); + } + } + /// Set credentials for HTTP request placeholder injection. pub fn with_credentials(mut self, credentials: HashMap) -> Self { self.credentials = credentials; @@ -615,9 +712,17 @@ impl WasmToolWrapper { } })?; + // Get typed interface — used for execute and error hints. + let tool_iface = instance.near_agent_tool(); + + // Determine effective schema for type coercion. + // Prefer the registration-time discovery schema when typed; otherwise + // try the WASM export transiently for this invocation only. + let effective_schema = self.schemas.effective_for_coercion(tool_iface, &mut store); + // Coerce string-encoded values to their schema-declared types. // LLMs frequently pass numeric values as strings (e.g. "5" instead of 5). - let params = coerce_params_to_schema(params, &self.schema); + let params = coerce_params_to_schema(params, &effective_schema); // Prepare the request let params_json = serde_json::to_string(¶ms) @@ -629,7 +734,6 @@ impl WasmToolWrapper { }; // Call execute using the generated typed interface - let tool_iface = instance.near_agent_tool(); let response = tool_iface.call_execute(&mut store, &request).map_err(|e| { let error_str = e.to_string(); if error_str.contains("out of fuel") { @@ -644,12 +748,13 @@ impl WasmToolWrapper { // Get logs from host state let logs = store.data_mut().host_state.take_logs(); - // Check for tool-level error — on failure, call the WASM module's - // description() and schema() exports so the LLM can retry with the - // correct parameters without us having to include the (large) schema - // in every request's tools array. + // Check for tool-level error — point the LLM to tool_info for the + // full schema instead of dumping ~3.5KB inline. if let Some(err) = response.error { - let hint = build_tool_hint(tool_iface, &mut store); + let hint = format!( + "Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.", + self.prepared.name + ); return Err(WasmError::ToolReturnedError { message: err, hint }); } @@ -658,47 +763,55 @@ impl WasmToolWrapper { } } -/// Maximum characters for the description portion of a tool hint. -const HINT_DESC_MAX: usize = 500; -/// Maximum characters for the schema portion of a tool hint. -const HINT_SCHEMA_MAX: usize = 3000; +/// Extract metadata (description + schema) from a WASM tool by briefly +/// instantiating it and calling its `description()` and `schema()` exports. +/// Analogous to MCP's `list_tools()` — discovers tool capabilities at load time. +/// +/// Falls back to generic description and permissive schema on failure. +pub(super) fn extract_wasm_metadata( + engine: &wasmtime::Engine, + component: &wasmtime::component::Component, + limits: &ResourceLimits, +) -> Result<(String, serde_json::Value), WasmError> { + let store_data = StoreData::new( + limits.memory_bytes, + Capabilities::default(), + HashMap::new(), + vec![], + ); + let mut store = Store::new(engine, store_data); -/// Call the WASM module's `description()` and `schema()` exports to build a -/// hint string. Returns an empty string if both calls fail or return empty. -/// Description is capped at [`HINT_DESC_MAX`] chars, schema at -/// [`HINT_SCHEMA_MAX`] chars. -fn build_tool_hint(tool_iface: &wit_tool::Guest, store: &mut Store) -> String { - let desc = tool_iface - .call_description(&mut *store) + // Configure fuel + epoch deadline so extraction can't hang + if let Err(e) = store.set_fuel(limits.fuel) { + tracing::debug!("Fuel not enabled for metadata extraction: {e}"); + } + store.epoch_deadline_trap(); + let ticks = (limits.timeout.as_millis() / EPOCH_TICK_INTERVAL.as_millis()).max(1) as u64; + store.set_epoch_deadline(ticks); + store.limiter(|data| &mut data.limiter); + + // Instantiate with minimal linker + let mut linker = Linker::new(engine); + WasmToolWrapper::add_host_functions(&mut linker)?; + let instance = SandboxedTool::instantiate(&mut store, component, &linker) + .map_err(|e| WasmError::InstantiationFailed(e.to_string()))?; + let tool_iface = instance.near_agent_tool(); + + // Extract description (fall back to generic) + let description = tool_iface + .call_description(&mut store) + .unwrap_or_else(|_| "WASM sandboxed tool".to_string()); + + // Extract and parse schema (fall back to permissive) + let schema = tool_iface + .call_schema(&mut store) .ok() - .unwrap_or_default(); - let schema = tool_iface.call_schema(&mut *store).ok().unwrap_or_default(); - if desc.is_empty() && schema.is_empty() { - return String::new(); - } - let mut hint = String::new(); - if !desc.is_empty() { - hint.push_str("Description: "); - if desc.len() > HINT_DESC_MAX { - let end = crate::util::floor_char_boundary(&desc, HINT_DESC_MAX); - hint.push_str(&desc[..end]); - hint.push('…'); - } else { - hint.push_str(&desc); - } - hint.push('\n'); - } - if !schema.is_empty() { - hint.push_str("Parameters schema: "); - if schema.len() > HINT_SCHEMA_MAX { - let end = crate::util::floor_char_boundary(&schema, HINT_SCHEMA_MAX); - hint.push_str(&schema[..end]); - hint.push('…'); - } else { - hint.push_str(&schema); - } - } - hint + .and_then(|s| serde_json::from_str::(&s).ok()) + .unwrap_or_else(|| { + serde_json::json!({"type": "object", "properties": {}, "additionalProperties": true}) + }); + + Ok((description, schema)) } #[async_trait] @@ -712,7 +825,11 @@ impl Tool for WasmToolWrapper { } fn parameters_schema(&self) -> serde_json::Value { - self.schema.clone() + self.schemas.advertised() + } + + fn discovery_schema(&self) -> serde_json::Value { + self.schemas.discovery() } async fn execute( @@ -749,7 +866,7 @@ impl Tool for WasmToolWrapper { let prepared = Arc::clone(&self.prepared); let capabilities = self.capabilities.clone(); let description = self.description.clone(); - let schema = self.schema.clone(); + let schemas = self.schemas.clone(); let credentials = self.credentials.clone(); // Execute in blocking task with timeout @@ -759,7 +876,7 @@ impl Tool for WasmToolWrapper { prepared, capabilities, description, - schema, + schemas, credentials, secrets_store: None, // Not needed in blocking task oauth_refresh: None, // Already used above for pre-refresh @@ -1232,6 +1349,7 @@ mod tests { TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET, test_secrets_store, }; + use crate::tools::tool::Tool; use crate::tools::wasm::capabilities::Capabilities; use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime}; @@ -1246,6 +1364,61 @@ mod tests { assert!(runtime.config().fuel_config.enabled); } + #[tokio::test] + async fn test_advertised_schema_stays_permissive_until_sidecar_override() { + let discovery_schema = serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "limit": { "type": "integer" } + }, + "required": ["query"] + }); + + let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::for_testing()).unwrap()); + let prepared = runtime + .prepare("search", b"\0asm\x0d\0\x01\0", None) + .await + .unwrap(); + let mut wrapper = + super::WasmToolWrapper::new(Arc::clone(&runtime), prepared, Capabilities::default()); + wrapper.schemas = super::WasmToolSchemas::new(discovery_schema.clone()); + wrapper.description = "Search documents".to_string(); + wrapper.append_schema_hint_if_permissive(); + + assert_eq!( + wrapper.parameters_schema(), + serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }) + ); + assert_eq!(wrapper.discovery_schema(), discovery_schema); + assert!(wrapper.description().contains("tool_info")); + + let wrapper = wrapper.with_schema(serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + })); + + assert_eq!( + wrapper.parameters_schema(), + serde_json::json!({ + "type": "object", + "properties": { + "query": { "type": "string" } + }, + "required": ["query"] + }) + ); + assert_eq!(wrapper.discovery_schema(), wrapper.parameters_schema()); + assert!(!wrapper.description().contains("tool_info")); + } + #[test] fn test_capabilities_default() { let caps = Capabilities::default(); @@ -1788,6 +1961,23 @@ mod tests { assert_eq!(result["count"], serde_json::json!("not-a-number")); } + /// Regression: permissive fallback schema (empty properties) must NOT coerce. + /// This documents the bug where WASM tools with no sidecar `parameters` field + /// got the permissive fallback, causing coercion to be a no-op and LLM-provided + /// string integers to reach the WASM tool un-coerced. + #[test] + fn test_coerce_noop_with_permissive_schema() { + let permissive = serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": true + }); + let params = serde_json::json!({"query": "test", "count": "10"}); + let result = super::coerce_params_to_schema(params, &permissive); + // With empty properties, no coercion happens — string stays string + assert_eq!(result["count"], serde_json::json!("10")); + } + /// Regression test: leak scan must run on raw headers (before credential /// injection), not after. If it ran post-injection, the host-injected /// Slack bot token (`xoxb-...`) would trigger a Block and reject the diff --git a/tests/e2e_builtin_tool_coverage.rs b/tests/e2e_builtin_tool_coverage.rs index 34cb35f7..f1ae3660 100644 --- a/tests/e2e_builtin_tool_coverage.rs +++ b/tests/e2e_builtin_tool_coverage.rs @@ -457,4 +457,90 @@ mod tests { rig.shutdown(); } + + // ----------------------------------------------------------------------- + // Test: tool_info_discovery (two-level detail) + // ----------------------------------------------------------------------- + // Verifies the tool_info built-in returns: + // - Default (no include_schema): name, description, parameter names array + // - With include_schema: true: adds full typed JSON Schema + + #[tokio::test] + async fn tool_info_discovery() { + let trace = LlmTrace::from_file(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/llm_traces/tools/tool_info_discovery.json" + )) + .expect("failed to load tool_info_discovery.json"); + + let rig = TestRigBuilder::new() + .with_trace(trace.clone()) + .with_auto_approve_tools(true) + .build() + .await; + + rig.send_message("What is the schema for the echo and time tools?") + .await; + let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await; + + rig.verify_trace_expects(&trace, &responses); + + // tool_info should have been called twice (echo + time), both succeeding. + let completed = rig.tool_calls_completed(); + let tool_info_calls: Vec<_> = completed.iter().filter(|(n, _)| n == "tool_info").collect(); + assert_eq!( + tool_info_calls.len(), + 2, + "Expected 2 tool_info calls, got {tool_info_calls:?}" + ); + assert!( + tool_info_calls.iter().all(|(_, ok)| *ok), + "All tool_info calls should succeed: {tool_info_calls:?}" + ); + + // Verify the results contain expected fields. + let results = rig.tool_results(); + let info_results: Vec<_> = results.iter().filter(|(n, _)| n == "tool_info").collect(); + + // First call was for "echo" (default, no include_schema) — result should + // contain "echo" and "parameters" as an array of names (not full schema). + let echo_result = info_results + .iter() + .find(|(_, preview)| preview.contains("echo")) + .expect("tool_info result should contain 'echo'"); + assert!( + echo_result.1.contains("message"), + "echo default result should list 'message' parameter name: {:?}", + echo_result.1 + ); + // Default mode should NOT include the full "schema" key + let echo_json: serde_json::Value = serde_json::from_str(&echo_result.1) + .expect("echo tool_info result should be valid JSON"); + assert!( + echo_json.get("schema").is_none(), + "Default tool_info should not include schema field: {:?}", + echo_result.1 + ); + + // Second call was for "time" with include_schema: true — result should + // contain "time", "schema" field with full object. + let time_result = info_results + .iter() + .find(|(_, preview)| preview.contains("time")) + .expect("tool_info result should contain 'time'"); + let time_json: serde_json::Value = serde_json::from_str(&time_result.1) + .expect("time tool_info result should be valid JSON"); + assert!( + time_json.get("schema").is_some(), + "include_schema: true should include schema field: {:?}", + time_result.1 + ); + assert!( + time_json["schema"]["properties"].is_object(), + "schema should have properties: {:?}", + time_result.1 + ); + + rig.shutdown(); + } } diff --git a/tests/fixtures/llm_traces/tools/tool_info_discovery.json b/tests/fixtures/llm_traces/tools/tool_info_discovery.json new file mode 100644 index 00000000..dc8746ad --- /dev/null +++ b/tests/fixtures/llm_traces/tools/tool_info_discovery.json @@ -0,0 +1,50 @@ +{ + "model_name": "test-tool-info-discovery", + "expects": { + "tools_used": ["tool_info"], + "all_tools_succeeded": true, + "min_responses": 1, + "tool_results_contain": { + "tool_info": "echo" + } + }, + "steps": [ + { + "request_hint": { "last_user_message_contains": "schema" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_tool_info_echo", + "name": "tool_info", + "arguments": { "name": "echo" } + } + ], + "input_tokens": 100, + "output_tokens": 20 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_tool_info_time", + "name": "tool_info", + "arguments": { "name": "time", "include_schema": true } + } + ], + "input_tokens": 200, + "output_tokens": 20 + } + }, + { + "response": { + "type": "text", + "content": "I found the info for both tools. The echo tool has a 'message' parameter. The time tool accepts an 'operation' parameter with options like 'now', 'parse', and 'diff'.", + "input_tokens": 400, + "output_tokens": 40 + } + } + ] +} diff --git a/tools-src/web-search/web-search-tool.capabilities.json b/tools-src/web-search/web-search-tool.capabilities.json index bc660aaf..9c2559ab 100644 --- a/tools-src/web-search/web-search-tool.capabilities.json +++ b/tools-src/web-search/web-search-tool.capabilities.json @@ -1,6 +1,41 @@ { "version": "0.2.0", "wit_version": "0.3.0", + "description": "Search the web using Brave Search. Returns titles, URLs, descriptions, and publication dates for matching web pages. Supports filtering by country, language, and freshness. Authentication is handled via the 'brave_api_key' secret injected by the host.", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to look up on the web" + }, + "count": { + "type": "integer", + "description": "Number of results to return (1-20, default 5)", + "minimum": 1, + "maximum": 20, + "default": 5 + }, + "country": { + "type": "string", + "description": "2-letter uppercase country code to bias results (e.g. 'US', 'DE', 'JP')" + }, + "search_lang": { + "type": "string", + "description": "2-letter lowercase language code for search results (e.g. 'en', 'de', 'fr')" + }, + "ui_lang": { + "type": "string", + "description": "Locale in language-region format (e.g. 'en-US', 'de-DE')" + }, + "freshness": { + "type": "string", + "description": "Filter by discovery time: 'pd' (past day), 'pw' (past week), 'pm' (past month), 'py' (past year), or date range 'YYYY-MM-DDtoYYYY-MM-DD'" + } + }, + "required": ["query"], + "additionalProperties": false + }, "capabilities": { "http": { "allowlist": [