mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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
This commit is contained in:
Henry Park
@@ -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));
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<ToolRegistry>,
|
||||
}
|
||||
|
||||
impl ToolInfoTool {
|
||||
pub fn new(registry: Weak<ToolRegistry>) -> 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<ToolOutput, ToolError> {
|
||||
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(_))));
|
||||
}
|
||||
}
|
||||
@@ -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<Self>` 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<Self>) {
|
||||
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:
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+6
-84
@@ -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<WasmError> 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<TrapCode>,
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
+26
-41
@@ -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<String, WasmError> {
|
||||
// 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<serde_json::Value, WasmError> {
|
||||
// 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")
|
||||
|
||||
+245
-55
@@ -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<String, String>,
|
||||
@@ -477,6 +477,79 @@ pub struct WasmToolWrapper {
|
||||
oauth_refresh: Option<OAuthRefreshConfig>,
|
||||
}
|
||||
|
||||
#[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<StoreData>,
|
||||
) -> 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::<serde_json::Value>(&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<PreparedModule>,
|
||||
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<String>) -> 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<String, String>) -> 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<StoreData>) -> 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::<serde_json::Value>(&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
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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": [
|
||||
|
||||
Reference in New Issue
Block a user