fix: add tool_info schema discovery for WASM tools#1086
fix: add tool_info schema discovery for WASM tools#1086henrypark133 merged 3 commits intostagingfrom
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
There was a problem hiding this comment.
Pull request overview
This PR introduces on-demand tool schema discovery via a new tool_info built-in tool, keeping the normal LLM tool advertisement compact (especially for large WASM tool schemas) while still enabling rich schema access for coercion and targeted retries.
Changes:
- Add
tool_infobuilt-in to fetch compact tool info by default and full typed JSON Schema on demand (include_schema: true). - Split WASM tool schemas into a compact advertised schema vs a full discovery/coercion schema, with per-turn hint dedup reset wired into the shared agentic loop lifecycle.
- Add regression/e2e coverage for tool discovery behavior, per-turn reset hook behavior, and advertised-vs-discovery schema handling.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tools-src/web-search/web-search-tool.capabilities.json | Adds explicit description + parameters schema in sidecar for a WASM tool. |
| tests/fixtures/llm_traces/tools/tool_info_discovery.json | Adds an LLM trace fixture exercising tool_info default vs full-schema mode. |
| tests/e2e_builtin_tool_coverage.rs | Adds an e2e coverage test validating tool_info discovery behavior/results. |
| src/worker/job.rs | Hooks notify_turn_start() into job loop turns via LoopDelegate::on_turn_start. |
| src/worker/container.rs | Hooks notify_turn_start() into container loop turns via LoopDelegate::on_turn_start. |
| src/tools/wasm/wrapper.rs | Splits advertised vs discovery schema, adds lazy coercion schema caching + per-turn hint dedup, and updates WASM error hinting to point to tool_info. |
| src/tools/wasm/runtime.rs | Switches registration-time metadata extraction to extract_wasm_metadata() with limits/timeouts and fallback behavior. |
| src/tools/wasm/mod.rs | Narrows public re-exports (removes trap-related exports). |
| src/tools/wasm/limits.rs | Removes unused limiter bookkeeping fields. |
| src/tools/wasm/error.rs | Removes unused trap structs/enums; updates hint formatting expectations in tests. |
| src/tools/tool.rs | Extends Tool trait with on_turn_start() and discovery_schema(). |
| src/tools/registry.rs | Adds notify_turn_start() and register_tool_info(); protects tool_info name. |
| src/tools/builtin/tool_info.rs | Implements the tool_info built-in tool + unit tests. |
| src/tools/builtin/mod.rs | Wires ToolInfoTool into the built-in module exports. |
| src/app.rs | Registers tool_info during app tool initialization. |
| src/agent/dispatcher.rs | Calls notify_turn_start() at the start of chat turns via delegate hook. |
| src/agent/agentic_loop.rs | Adds LoopDelegate::on_turn_start() and calls it once per run; adds tests for once-per-loop behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
zmanian
left a comment
There was a problem hiding this comment.
Reviewed the full diff. This is well-structured and solves a real problem (reducing token usage from full WASM schema dumps).
Looks good:
tool_infocorrectly added toPROTECTED_TOOL_NAMES- Two-tier schema system (advertised vs discovery) is clean
- Good test coverage: unit tests for both modes,
WasmToolSchemasbehavior, and E2E trace test - Dead code cleanup (
TrapInfo,TrapCode, unused fields) is welcome - Backward-compatible
discovery_schema()trait method with sensible default
Minor suggestions (non-blocking):
WasmToolSchemas::advertised()clones the schema on every call viaparameters_schema(), which runs on every LLM request. Consider returning&serde_json::Valueinstead.- The hint text manipulation (
append_schema_hint_if_permissive/strip_schema_hint) usingcontains/find/truncateis fragile. A separateschema_hint: Option<String>field composed at display time would be more robust. ToolInfoToolholdsArc<ToolRegistry>while the registry holdsToolInfoTool-- circularArcreference. Fine for app-lifetime objects, butWeak<ToolRegistry>would be cleaner if this ever gets refactored.
LGTM
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| /// 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"); | ||
| } |
|
|
||
| // Core types | ||
| pub use error::{TrapCode, TrapInfo, WasmError}; | ||
| pub use error::WasmError; |
| 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(); |
| /// 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)"); | ||
| } |
zmanian
left a comment
There was a problem hiding this comment.
Design Validations
The core architecture here is solid and well-motivated:
-
Two-tier schema system (advertised vs discovery) is the right approach. Keeping the tools array compact while exposing full schemas on demand directly addresses the prompt-size regression. The
WasmToolSchemasstruct makes the two roles explicit and immutable. -
discovery_schema()trait method with the default delegation toparameters_schema()is backward-compatible and clean. Non-WASM tools don't need to care about it. -
tool_infoinPROTECTED_TOOL_NAMESis correct -- prevents WASM tools from shadowing it. -
Removing per-turn hint dedup is the right call. Repeated hints on repeated failures give the LLM useful signal about the failure pattern. The old once-per-turn mechanism was hiding information.
-
Dead code cleanup (TrapInfo, TrapCode, unused limiter fields) is welcome. Confirmed these types have zero consumers outside their own module and tests.
-
E2E test with LLM trace fixture covers both detail levels and verifies the tool is wired end-to-end.
Actionable Feedback
1. Description hint mutation is fragile (blocking)
append_schema_hint_if_permissive() / strip_schema_hint() mutate the stored description string using contains / find / truncate. This couples presentation concerns to stored state and breaks if a tool description naturally contains "tool_info".
Suggestion: Store a schema_hint_needed: bool field on WasmToolSchemas (it already knows whether advertised is permissive). Then compose the hint in the Tool::schema() method which returns an owned ToolSchema, or in description() if you change the return type. The raw description stays clean and the hint is purely a presentation concern. This also addresses Copilot's comment about the hint text not mentioning include_schema: true -- you can fix the hint text in one place.
2. extract_wasm_metadata does redundant work (non-blocking, but worth addressing)
The new extract_wasm_metadata instantiates each WASM module during prepare() to call description() and schema() exports. But when a sidecar capabilities.json exists (which is the common case), the loader immediately overrides both via with_description() / with_schema(), throwing away the extraction results.
Suggestion: Either:
- Pass optional description/schema into
prepare()so it can skip WASM instantiation when the sidecar already provides them - Or make extraction lazy -- only instantiate when
discovery_schema()is first called and no sidecar override was set
For registries with many WASM tools, this avoids doubling the instantiation cost at startup.
3. effective_for_coercion transient fallback needs safety documentation (non-blocking)
When the discovery schema is permissive, effective_for_coercion calls tool_iface.call_schema(&mut store) on the already-running WASM instance between store setup and call_execute(). This has two risks:
- State contamination: The WASM module's linear memory is already initialized. If
schema()reads or writes mutable state, calling it mid-execution could corrupt the environment beforecall_execute(). - Inconsistency: If
schema()returns different results depending on internal state (e.g., post-initialization), the coercion schema could differ from the load-time extraction.
In practice, well-behaved tools return a static string, so this is likely safe. But please either:
- Add a comment documenting this assumption ("schema() is assumed to be a pure function with no side effects")
- Or preferably, use the load-time extracted schema from
PreparedModule.schemainstead of re-calling the export.PreparedModulealready holds it and is available viaself.prepared.
4. register_tool_info is easy to forget (non-blocking)
Copilot flagged this and it's valid. register_tool_info() is a separate call from register_builtin_tools(), which means test harnesses and other call sites can forget it. Then WASM tool descriptions say "call tool_info" but the tool isn't registered.
Consider either folding it into a register_builtin_tools(self: &Arc<Self>) that takes Arc<Self>, or adding a register_all_builtins(self: &Arc<Self>) convenience method.
5. Minor: WasmToolSchemas::advertised() clones on every call
parameters_schema() is called on every LLM request for every tool. advertised() does a full serde_json::Value clone each time. Consider returning &serde_json::Value or wrapping in Arc<serde_json::Value>.
…rcion safety Two fixes from the review of #1086 (tool_info schema discovery): 1. Replace fragile description string mutation (append_schema_hint_if_permissive / strip_schema_hint) with composition at display time. The raw description stays clean; the tool_info hint is composed in the Tool::schema() override only when the advertised schema is permissive. This also includes the tool name and `include_schema: true` in the hint for better LLM guidance. 2. Make effective_for_coercion use the load-time extracted schema from PreparedModule instead of re-calling the WASM schema() export on the already-running instance mid-execution. This avoids potential state contamination from calling schema() after linear memory is initialized for execution. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…rcion safety (#1092) Two fixes from the review of #1086 (tool_info schema discovery): 1. Replace fragile description string mutation (append_schema_hint_if_permissive / strip_schema_hint) with composition at display time. The raw description stays clean; the tool_info hint is composed in the Tool::schema() override only when the advertised schema is permissive. This also includes the tool name and `include_schema: true` in the hint for better LLM guidance. 2. Make effective_for_coercion use the load-time extracted schema from PreparedModule instead of re-calling the WASM schema() export on the already-running instance mid-execution. This avoids potential state contamination from calling schema() after linear memory is initialized for execution. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
* fix: add tool_info schema discovery for WASM tools * refactor: simplify WASM schema and hint state * refactor: store tool_info registry reference as Weak
…nd coercion safety (nearai#1092) Two fixes from the review of nearai#1086 (tool_info schema discovery): 1. Replace fragile description string mutation (append_schema_hint_if_permissive / strip_schema_hint) with composition at display time. The raw description stays clean; the tool_info hint is composed in the Tool::schema() override only when the advertised schema is permissive. This also includes the tool name and `include_schema: true` in the hint for better LLM guidance. 2. Make effective_for_coercion use the load-time extracted schema from PreparedModule instead of re-calling the WASM schema() export on the already-running instance mid-execution. This avoids potential state contamination from calling schema() after linear memory is initialized for execution. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Summary
This PR adds an on-demand schema discovery path for tools and fixes the prompt-size regression for WASM tools by separating the schema sent in the normal
toolsarray from the richer schema exposed throughtool_info.The final design is intentionally simpler than the intermediate version on this branch:
tool_infoWhat Changed
tool_infotool that returns:include_schema: trueTooltrait withdiscovery_schema()so tools can expose richer discovery data than the compactparameters_schema()used in the normaltoolsarrayparameters_schema()returns the advertised schema used in the main promptdiscovery_schema()returns the richer registration-time discovery schema used bytool_infoWasmToolSchemashelper to hold the two immutable schema roles:advertisedfor the normal tools arraydiscoveryfortool_infoand typed coercion when availableschema()export for that invocation onlytool_info(name: ..., include_schema: true)retry hint on WASM tool-level errorsWhy This Approach
This keeps the common path cheap while still letting the model discover precise schemas when it needs them:
toolsarray stays small and stabletool_infoon demand for the full schemaThe simplification also removes branch-specific state that no longer pays for itself:
hint_sentflagReview Fixes Included
This PR addresses the review issues raised on the branch:
tool_inforetry hint behavior consistent across chat, job, and container flows by removing the once-per-turn dedup mechanism entirely and always emitting the short hintTesting
cargo fmtcargo clippy --all-targets --all-features -- -D warningscargo test test_advertised_schema_stays_permissive_until_sidecar_override -- --nocapturecargo test test_tool_intent_nudge_fires_and_caps -- --nocapture