Skip to content

feat(wasm): lazy schema injection on WASM tool errors #638

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
henrypark133 merged 2 commits into from
Mar 9, 2026
Merged

feat(wasm): lazy schema injection on WASM tool errors #638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8b82c87
feat(wasm): lazy schema injection on WASM tool errors
henrypark133 Mar 6, 2026
225d58e
fix: use floor_char_boundary for UTF-8 safe truncation in tool hints
henrypark133 Mar 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions src/tools/wasm/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,15 @@ pub enum WasmError {
Timeout(std::time::Duration),

/// Component returned an error response.
#[error("Tool error: {0}")]
ToolReturnedError(String),
/// When `hint` is non-empty it carries the tool's description and parameter
/// schema so the LLM can retry with correct arguments.
#[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).
hint: String,
},
Comment on lines +71 to +79

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The #[error(...)] string builds the hint suffix with an inline if ... { format!(...) } expression, which makes the derive attribute hard to read/maintain. Consider storing the fully-formatted hint (including any leading newlines/labels) in the hint field when constructing the error (or moving the suffix formatting into a small helper) so the error message stays a simple "Tool error: {message}{hint}".

Copilot uses AI. Check for mistakes.

@henrypark133 henrypark133 Mar 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The inline conditional is a standard thiserror pattern — it keeps the formatting co-located with the variant definition. Moving it elsewhere adds indirection without functional benefit. Leaving as-is.


/// Invalid JSON in tool response.
#[error("Invalid response JSON: {0}")]
Expand Down Expand Up @@ -195,4 +202,28 @@ mod tests {
_ => panic!("Expected Sandbox variant"),
}
}

#[test]
fn test_tool_returned_error_without_hint() {
let err = WasmError::ToolReturnedError {
message: "unknown action: foobar".to_string(),
hint: String::new(),
};
let display = err.to_string();
assert!(display.contains("unknown action: foobar"));
assert!(!display.contains("Tool usage hint"));
}

#[test]
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(),
};
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"));
}
}
51 changes: 49 additions & 2 deletions src/tools/wasm/wrapper.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -633,16 +633,63 @@ impl WasmToolWrapper {
// Get logs from host state
let logs = store.data_mut().host_state.take_logs();

// Check for tool-level error
// 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.
if let Some(err) = response.error {
return Err(WasmError::ToolReturnedError(err));
let hint = build_tool_hint(tool_iface, &mut store);
return Err(WasmError::ToolReturnedError { message: err, hint });
}

// Return result (or empty string if none)
Ok((response.output.unwrap_or_default(), logs))
}
}

/// 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;

/// 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)
.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('…');
Comment on lines +682 to +685

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

schema.len()/&schema[..HINT_SCHEMA_MAX] has the same UTF-8 boundary panic risk as the description truncation and also enforces a byte cap (despite docs saying "chars"). Please truncate safely at a char boundary (or rename to *_MAX_BYTES) so this can’t panic on non-ASCII schemas.

Copilot uses AI. Check for mistakes.

@henrypark133 henrypark133 Mar 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same fix applied — 225d58e.

} else {
hint.push_str(&schema);
Comment on lines +655 to +687

Copilot AI Mar 6, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

desc.len()/&desc[..HINT_DESC_MAX] treats the limit as bytes and can panic if HINT_DESC_MAX lands mid–UTF-8 codepoint (non-ASCII descriptions). Consider truncating on a char boundary (e.g., compute the byte offset via char_indices().nth(HINT_DESC_MAX) or walk back with is_char_boundary) and update the comment/constants to clarify whether the cap is bytes vs chars.

Suggested change
/// 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)
.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 {
hint.push_str(&desc[..HINT_DESC_MAX]);
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 {
hint.push_str(&schema[..HINT_SCHEMA_MAX]);
hint.push('…');
} else {
hint.push_str(&schema);
/// Truncate a UTF-8 string to at most `max_chars` Unicode scalar values.
/// Returns a slice of the original string and a flag indicating whether
/// truncation occurred.
fn truncate_utf8_to_char_boundary<'a>(s: &'a str, max_chars: usize) -> (&'a str, bool) {
if max_chars == 0 {
return ("", !s.is_empty());
}
let mut char_count = 0usize;
let mut byte_idx = s.len();
for (i, _) in s.char_indices() {
if char_count == max_chars {
byte_idx = i;
break;
}
char_count += 1;
}
if char_count <= max_chars && byte_idx == s.len() {
return (s, false);
}
(&s[..byte_idx], true)
}
/// 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`] characters, schema at
/// [`HINT_SCHEMA_MAX`] characters.
fn build_tool_hint(tool_iface: &wit_tool::Guest, store: &mut Store<StoreData>) -> String {
let desc = tool_iface
.call_description(&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: ");
let (desc_truncated, was_truncated) =
truncate_utf8_to_char_boundary(&desc, HINT_DESC_MAX);
hint.push_str(desc_truncated);
if was_truncated {
hint.push('…');
}
hint.push('\n');
}
if !schema.is_empty() {
hint.push_str("Parameters schema: ");
let (schema_truncated, was_truncated) =
truncate_utf8_to_char_boundary(&schema, HINT_SCHEMA_MAX);
hint.push_str(schema_truncated);
if was_truncated {
hint.push('…');

Copilot uses AI. Check for mistakes.

@henrypark133 henrypark133 Mar 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 225d58e — using crate::util::floor_char_boundary() for both description and schema truncation.

}
}
hint
Comment on lines +660 to +690

@gemini-code-assist gemini-code-assist Bot Mar 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-high high

The build_tool_hint function is vulnerable to a high-severity prompt injection risk. It directly includes untrusted output from WASM module's description() and schema() exports into hints presented to the LLM. A malicious WASM tool could craft these strings to inject harmful instructions, potentially leading to unauthorized actions. To mitigate this, sanitize the output or use structured formats, as per Rule: 'Tool output previews sent to trusted local endpoints (e.g., TUI, web gateway) should not be sanitized. Sanitization should only be applied to data paths sent to external services, such as an LLM.'. Additionally, the current string truncation logic using &desc[..HINT_DESC_MAX] and &schema[..HINT_SCHEMA_MAX] is unsafe and can cause a panic with multi-byte UTF-8 characters. This violates Rule: 'When truncating a UTF-8 string, use character-aware methods (e.g., char_indices() or iterating over chars()) to avoid panics caused by slicing in the middle of a multi-byte character.' and Rule: 'When truncating a UTF-8 string at a byte boundary, walk backwards from the desired length until a valid character boundary is found using is_char_boundary to prevent panics.'. This truncation logic is also duplicated and should be refactored into a safe helper function for better maintainability and correctness, also aligning with Rule: 'Always truncate tool output for previews or status updates to a reasonable maximum length. This prevents excessive memory/bandwidth usage and reduces the risk of leaking sensitive information.'.

    let mut hint = String::new();

    fn append_truncated(buf: &mut String, s: &str, max_len: usize) {
        if s.len() > max_len {
            let mut end = max_len;
            while !s.is_char_boundary(end) {
                end -= 1;
            }
            buf.push_str(&s[..end]);
            buf.push('…');
        } else {
            buf.push_str(s);
        }
    }

    if !desc.is_empty() {
        hint.push_str("Description: ");
        append_truncated(&mut hint, &desc, HINT_DESC_MAX);
        hint.push('\n');
    }
    if !schema.is_empty() {
        hint.push_str("Parameters schema: ");
        append_truncated(&mut hint, &schema, HINT_SCHEMA_MAX);
    }
    hint
References
  1. Tool output previews sent to trusted local endpoints (e.g., TUI, web gateway) should not be sanitized. Sanitization should only be applied to data paths sent to external services, such as an LLM.
  2. When truncating a UTF-8 string, use character-aware methods (e.g., char_indices() or iterating over chars()) to avoid panics caused by slicing in the middle of a multi-byte character.
  3. When truncating a UTF-8 string at a byte boundary, walk backwards from the desired length until a valid character boundary is found using is_char_boundary to prevent panics.
  4. Always truncate tool output for previews or status updates to a reasonable maximum length. This prevents excessive memory/bandwidth usage and reduces the risk of leaking sensitive information.

@henrypark133 henrypark133 Mar 7, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two points here:

Prompt injection: This isn't new attack surface. The error path at worker.rs:986 already sends format!("Error: {}", e) as a tool result to the LLM — same as every other tool error (built-in or WASM). The WASM tool's execute() error message already flows unsanitized to the LLM on the error path. The hint (description/schema) is the tool's own metadata, no different from what would be in the tools array if we included it there.

UTF-8 panic: Valid catch. Fixed in 225d58e — now uses crate::util::floor_char_boundary(), which is the existing polyfill used in 6+ places across the codebase for exactly this purpose.

}

#[async_trait]
impl Tool for WasmToolWrapper {
fn name(&self) -> &str {
Expand Down