feat(wasm): lazy schema injection on WASM tool errors - #638
Conversation
When a WASM tool returns an error (ToolReturnedError), call the module's description() and schema() WIT exports and append them as a hint in the error message. This lets the LLM retry with correct parameters without us including large schemas in every request's tools array. - Change ToolReturnedError from tuple to struct variant with hint field - Add build_tool_hint() that calls WASM description()/schema() exports - Cap description at 500 chars, schema at 3000 chars to limit context - Hint flows automatically through Display → ToolError → ChatMessage Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
scope: tool/wasm
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly improves the interaction between LLMs and WASM tools by implementing a "lazy schema injection" mechanism. Instead of pre-loading potentially large tool schemas, the system now dynamically extracts and provides tool descriptions and parameter schemas as hints only when a WASM tool encounters an error. This approach helps the LLM correct its subsequent calls, reduces token usage by avoiding constant schema inclusion, and makes WASM tool integration more robust and cost-effective. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a useful feature for providing lazy hints to the LLM when a WASM tool fails, aiming to improve its ability to self-correct. However, the current implementation presents a high-severity prompt injection risk due to untrusted WASM tool exports being directly used as hints to the LLM, which requires sanitization for external services like LLMs. Additionally, the unsafe string truncation in the build_tool_hint function poses a denial-of-service risk with multi-byte UTF-8 characters, violating best practices for character-aware truncation. My review includes a comprehensive suggestion to address these issues by using a safe, character-aware truncation method and refactoring duplicated logic for better maintainability and security.
| 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); | ||
| } | ||
| } | ||
| hint |
There was a problem hiding this comment.
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);
}
hintReferences
- 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.
- When truncating a UTF-8 string, use character-aware methods (e.g.,
char_indices()or iterating overchars()) to avoid panics caused by slicing in the middle of a multi-byte character. - 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_boundaryto prevent panics. - 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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Pull request overview
This PR adds lazy tool-schema/description injection for WASM tool failures: when a WASM tool returns an error, the host calls the module’s description() and schema() exports and appends them as a hint in the error string so the LLM can retry with correct arguments without including large schemas in every request.
Changes:
- Inject a “tool usage hint” (description + parameters schema) into
ToolReturnedErrorwhen a WASM tool returns an error. - Change
WasmError::ToolReturnedErrorfrom a tuple variant to a struct variant{ message, hint }and update its Display formatting. - Add tests validating that the hint is omitted when empty and included when present.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/tools/wasm/wrapper.rs |
Builds and attaches a description/schema hint on tool-level error responses. |
src/tools/wasm/error.rs |
Updates the error variant to carry an optional hint and adds Display/tests for the new format. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| /// 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); |
There was a problem hiding this comment.
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.
| /// 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('…'); |
There was a problem hiding this comment.
Fixed in 225d58e — using crate::util::floor_char_boundary() for both description and schema truncation.
| if schema.len() > HINT_SCHEMA_MAX { | ||
| hint.push_str(&schema[..HINT_SCHEMA_MAX]); | ||
| hint.push('…'); |
There was a problem hiding this comment.
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.
| /// 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, | ||
| }, |
There was a problem hiding this comment.
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}".
There was a problem hiding this comment.
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.
Use existing crate::util::floor_char_boundary() to avoid panicking when truncation lands mid-multibyte character. Addresses review feedback on PR #638. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* feat(wasm): lazy schema injection on WASM tool errors When a WASM tool returns an error (ToolReturnedError), call the module's description() and schema() WIT exports and append them as a hint in the error message. This lets the LLM retry with correct parameters without us including large schemas in every request's tools array. - Change ToolReturnedError from tuple to struct variant with hint field - Add build_tool_hint() that calls WASM description()/schema() exports - Cap description at 500 chars, schema at 3000 chars to limit context - Hint flows automatically through Display → ToolError → ChatMessage Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use floor_char_boundary for UTF-8 safe truncation in tool hints Use existing crate::util::floor_char_boundary() to avoid panicking when truncation lands mid-multibyte character. Addresses review feedback on PR nearai#638. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* feat(wasm): lazy schema injection on WASM tool errors When a WASM tool returns an error (ToolReturnedError), call the module's description() and schema() WIT exports and append them as a hint in the error message. This lets the LLM retry with correct parameters without us including large schemas in every request's tools array. - Change ToolReturnedError from tuple to struct variant with hint field - Add build_tool_hint() that calls WASM description()/schema() exports - Cap description at 500 chars, schema at 3000 chars to limit context - Hint flows automatically through Display → ToolError → ChatMessage Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix: use floor_char_boundary for UTF-8 safe truncation in tool hints Use existing crate::util::floor_char_boundary() to avoid panicking when truncation lands mid-multibyte character. Addresses review feedback on PR nearai#638. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Summary
description()andschema()WIT exports and inject them as a hint in the error responsetoolsarrayProblem
WASM tools define
description()andschema()exports in their WIT interface, butextract_tool_description()inruntime.rsis a stub that always returns"WASM sandboxed tool"with an empty schema. The LLM has to guess parameter names and gets it wrong, wasting tokens on retries.Including full schemas in every request would be expensive (some Google API schemas are huge). This PR takes the lazy injection approach: only inject when the tool returns an error.
Changes
src/tools/wasm/error.rs—ToolReturnedErrorchanged from(String)to{ message, hint }struct variant. Display appends hint when non-empty.src/tools/wasm/wrapper.rs—build_tool_hint()calls the WASM module'sdescription()/schema()exports on the already-instantiated module whenexecute()returns an error.Test plan
cargo fmt— cleancargo clippy --all --all-features— zero warningstest_tool_returned_error_without_hint— verifies hint is omitted when emptytest_tool_returned_error_with_hint— verifies hint appears in Display output[skip-regression-check]
Generated with Claude Code