Fix schema-guided tool parameter coercion (#1143)

* Fix schema-guided tool parameter coercion

* Fix CI checks for coercion regression tests

* Finish panic-scan annotations

* Avoid redundant worker param preparation

* Keep panic-scan annotations rustfmt-stable

* Handle nullable WASM schema review feedback

* Address param coercion review notes
This commit is contained in:
Henry Park
2026-03-14 16:27:18 -07:00
committed by GitHub
parent fda5160940
commit c79754df28
9 changed files with 1025 additions and 193 deletions
+10 -4
View File
@@ -32,7 +32,9 @@ use crate::llm::{
ChatMessage, CompletionRequest, FinishReason, LlmProvider, ToolCall, ToolCompletionRequest,
};
use crate::safety::SafetyLayer;
use crate::tools::{ApprovalContext, ApprovalRequirement, ToolError, ToolRegistry};
use crate::tools::{
ApprovalContext, ApprovalRequirement, ToolError, ToolRegistry, prepare_tool_params,
};
use crate::workspace::Workspace;
enum EventMatcher {
@@ -1118,13 +1120,14 @@ async fn execute_routine_tool(
.get(&tc.name)
.await
.ok_or_else(|| format!("Tool '{}' not found", tc.name))?;
let normalized_params = prepare_tool_params(tool.as_ref(), &tc.arguments);
// Check approval requirement: only allow Never tools in lightweight routines.
// UnlessAutoApproved and Always tools are blocked to prevent prompt injection attacks.
// Lightweight routines can be triggered by external events and may process untrusted data,
// making them vulnerable to prompt injection that could trick the LLM into calling
// sensitive tools. Blocking these tools entirely is the safest approach.
match tool.requires_approval(&tc.arguments) {
match tool.requires_approval(&normalized_params) {
ApprovalRequirement::Never => {}
ApprovalRequirement::UnlessAutoApproved | ApprovalRequirement::Always => {
return Err(format!(
@@ -1136,7 +1139,10 @@ async fn execute_routine_tool(
}
// Validate tool parameters
let validation = ctx.safety.validator().validate_tool_params(&tc.arguments);
let validation = ctx
.safety
.validator()
.validate_tool_params(&normalized_params);
if !validation.is_valid {
let details = validation
.errors
@@ -1151,7 +1157,7 @@ async fn execute_routine_tool(
let timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(timeout, async {
tool.execute(tc.arguments.clone(), job_ctx).await
tool.execute(normalized_params.clone(), job_ctx).await
})
.await;
let elapsed = start.elapsed();
+84 -3
View File
@@ -17,7 +17,7 @@ use crate::error::{Error, JobError};
use crate::hooks::HookRegistry;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::tools::{ApprovalContext, ToolRegistry};
use crate::tools::{ApprovalContext, ToolRegistry, prepare_tool_params};
use crate::worker::job::{Worker, WorkerDeps};
/// Message to send to a worker.
@@ -511,8 +511,10 @@ impl Scheduler {
.into());
}
let normalized_params = prepare_tool_params(tool.as_ref(), &params);
// Scheduler-specific approval check
let requirement = tool.requires_approval(&params);
let requirement = tool.requires_approval(&normalized_params);
let blocked =
ApprovalContext::is_blocked_or_default(&approval_context, tool_name, requirement);
if blocked {
@@ -524,7 +526,11 @@ impl Scheduler {
// Delegate to shared tool execution pipeline
let output_str = crate::tools::execute::execute_tool_with_safety(
&tools, &safety, tool_name, &params, &job_ctx,
&tools,
&safety,
tool_name,
&normalized_params,
&job_ctx,
)
.await?;
@@ -1064,4 +1070,79 @@ mod tests {
"hard_gate should pass with explicit permission"
);
}
struct NormalizedApprovalTool;
#[async_trait::async_trait]
impl Tool for NormalizedApprovalTool {
fn name(&self) -> &str {
"normalized_gate"
}
fn description(&self) -> &str {
"approval depends on normalized params"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"safe": { "type": "boolean" }
}
})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput::text(
"normalized_ok",
std::time::Instant::now().elapsed(),
))
}
fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement {
if params.get("safe").and_then(|v| v.as_bool()) == Some(true) {
ApprovalRequirement::Never
} else {
ApprovalRequirement::Always
}
}
fn requires_sanitization(&self) -> bool {
false
}
}
#[tokio::test]
async fn test_execute_tool_task_normalizes_params_before_approval() {
let registry = ToolRegistry::new();
registry.register(Arc::new(NormalizedApprovalTool)).await;
let cm = Arc::new(ContextManager::new(5));
let job_id = cm.create_job("test", "normalized approval").await.unwrap(); // safety: test-only setup
cm.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None))
.await
.unwrap() // safety: test-only setup
.unwrap(); // safety: test-only setup
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
}));
let result = Scheduler::execute_tool_task(
Arc::new(registry),
cm,
safety,
None,
job_id,
"normalized_gate",
serde_json::json!({"safe": "true"}),
)
.await;
#[rustfmt::skip]
assert!( // safety: test-only assertion
result.is_ok(),
"stringified boolean should normalize before approval: {result:?}"
);
}
}
+3 -2
View File
@@ -43,8 +43,8 @@ use crate::error::ToolError as AgentToolError;
use crate::llm::{
ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolDefinition,
};
use crate::tools::ToolRegistry;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput};
use crate::tools::{ToolRegistry, prepare_tool_params};
/// Requirement specification for building software.
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -776,10 +776,11 @@ Create alongside the .wasm file to grant capabilities:
self.tools.get(tool_name).await.ok_or_else(|| {
ToolError::ExecutionFailed(format!("Tool not found: {}", tool_name))
})?;
let normalized_params = prepare_tool_params(tool.as_ref(), params);
// Execute with a dummy context (build tools don't need job context)
let ctx = JobContext::default();
tool.execute(params.clone(), &ctx).await
tool.execute(normalized_params, &ctx).await
}
/// Find the build artifact based on project type.
+367
View File
@@ -0,0 +1,367 @@
pub(crate) fn prepare_tool_params(
tool: &dyn crate::tools::tool::Tool,
params: &serde_json::Value,
) -> serde_json::Value {
prepare_params_for_schema(params, &tool.discovery_schema())
}
pub(crate) fn prepare_params_for_schema(
params: &serde_json::Value,
schema: &serde_json::Value,
) -> serde_json::Value {
coerce_value(params, schema)
}
fn coerce_value(value: &serde_json::Value, schema: &serde_json::Value) -> serde_json::Value {
// This coercer intentionally handles the concrete schema shapes we expose in
// discovery today. It does not resolve combinators like anyOf/oneOf/allOf or
// references via $ref; those schemas pass through unchanged unless they also
// advertise a directly coercible type/property shape.
if value.is_null() {
return value.clone();
}
if let Some(s) = value.as_str() {
return coerce_string_value(s, schema).unwrap_or_else(|| value.clone());
}
if let Some(items) = value.as_array() {
if !schema_allows_type(schema, "array") {
return value.clone();
}
let Some(item_schema) = schema.get("items") else {
return value.clone();
};
return serde_json::Value::Array(
items
.iter()
.map(|item| coerce_value(item, item_schema))
.collect(),
);
}
if let Some(obj) = value.as_object() {
if !schema_allows_type(schema, "object") {
return value.clone();
}
let properties = schema.get("properties").and_then(|p| p.as_object());
let additional_schema = schema.get("additionalProperties").filter(|v| v.is_object());
let mut coerced = obj.clone();
for (key, current) in &mut coerced {
if let Some(prop_schema) = properties.and_then(|props| props.get(key)) {
*current = coerce_value(current, prop_schema);
continue;
}
if let Some(additional_schema) = additional_schema {
*current = coerce_value(current, additional_schema);
}
}
return serde_json::Value::Object(coerced);
}
value.clone()
}
fn coerce_string_value(s: &str, schema: &serde_json::Value) -> Option<serde_json::Value> {
if schema_allows_type(schema, "string") {
return None;
}
if schema_allows_type(schema, "integer")
&& let Ok(v) = s.parse::<i64>()
{
return Some(serde_json::Value::from(v));
}
if schema_allows_type(schema, "number")
&& let Ok(v) = s.parse::<f64>()
{
return Some(serde_json::Value::from(v));
}
if schema_allows_type(schema, "boolean") {
match s.to_lowercase().as_str() {
"true" => return Some(serde_json::json!(true)),
"false" => return Some(serde_json::json!(false)),
_ => {}
}
}
if schema_allows_type(schema, "array") || schema_allows_type(schema, "object") {
let parsed = serde_json::from_str::<serde_json::Value>(s).ok()?;
let matches_schema = match &parsed {
serde_json::Value::Array(_) => schema_allows_type(schema, "array"),
serde_json::Value::Object(_) => schema_allows_type(schema, "object"),
_ => false,
};
if matches_schema {
return Some(coerce_value(&parsed, schema));
}
}
None
}
fn schema_allows_type(schema: &serde_json::Value, expected: &str) -> bool {
match schema.get("type") {
Some(serde_json::Value::String(t)) => t == expected,
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
_ => match expected {
"object" => schema
.get("properties")
.and_then(|p| p.as_object())
.is_some(),
"array" => schema.get("items").is_some(),
_ => false,
},
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use async_trait::async_trait;
use super::*;
use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
struct StubTool {
schema: serde_json::Value,
}
#[async_trait]
impl Tool for StubTool {
fn name(&self) -> &str {
"stub"
}
fn description(&self) -> &str {
"stub"
}
fn parameters_schema(&self) -> serde_json::Value {
self.schema.clone()
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput::success(params, Duration::from_millis(1)))
}
}
#[test]
fn coerces_scalar_strings() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"count": { "type": "number" },
"limit": { "type": "integer" },
"enabled": { "type": "boolean" }
}
});
let params = serde_json::json!({
"count": "5",
"limit": "10",
"enabled": "TRUE"
});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["count"], serde_json::json!(5.0)); // safety: test-only assertion
assert_eq!(result["limit"], serde_json::json!(10)); // safety: test-only assertion
assert_eq!(result["enabled"], serde_json::json!(true)); // safety: test-only assertion
}
#[test]
fn coerces_stringified_array_and_recurses_into_items() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"values": {
"type": "array",
"items": {
"type": "array",
"items": { "type": "integer" }
}
}
}
});
let params = serde_json::json!({
"values": "[[\"1\", \"2\"], [\"3\", 4]]"
});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["values"], serde_json::json!([[1, 2], [3, 4]])); // safety: test-only assertion
}
#[test]
fn coerces_stringified_object_and_recurses_into_properties() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"request": {
"type": "object",
"properties": {
"start_index": { "type": "integer" },
"enabled": { "type": ["boolean", "null"] }
}
}
}
});
let params = serde_json::json!({
"request": "{\"start_index\":\"12\",\"enabled\":\"false\"}"
});
let result = prepare_params_for_schema(&params, &schema);
#[rustfmt::skip]
assert_eq!( // safety: test-only assertion
result["request"],
serde_json::json!({"start_index": 12, "enabled": false})
);
}
#[test]
fn coerces_nullable_stringified_arrays() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"requests": {
"type": ["array", "null"],
"items": {
"type": "object",
"properties": {
"enabled": { "type": "boolean" }
}
}
}
}
});
let params = serde_json::json!({
"requests": "[{\"enabled\":\"true\"}]"
});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["requests"], serde_json::json!([{ "enabled": true }])); // safety: test-only assertion
}
#[test]
fn coerces_typed_additional_properties() {
let schema = serde_json::json!({
"type": "object",
"additionalProperties": {
"type": "object",
"properties": {
"count": { "type": "integer" },
"enabled": { "type": "boolean" }
}
}
});
let params = serde_json::json!({
"alpha": "{\"count\":\"5\",\"enabled\":\"false\"}",
"beta": { "count": "7", "enabled": "true" }
});
let result = prepare_params_for_schema(&params, &schema);
#[rustfmt::skip]
assert_eq!( // safety: test-only assertion
result,
serde_json::json!({
"alpha": { "count": 5, "enabled": false },
"beta": { "count": 7, "enabled": true }
})
);
}
#[test]
fn leaves_invalid_json_strings_unchanged() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"requests": {
"type": "array",
"items": { "type": "object" }
}
}
});
let params = serde_json::json!({
"requests": "[{\"oops\":]"
});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["requests"], serde_json::json!("[{\"oops\":]")); // safety: test-only assertion
}
#[test]
fn leaves_string_when_schema_allows_string() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"value": { "type": ["string", "object"] }
}
});
let params = serde_json::json!({
"value": "{\"mode\":\"raw\"}"
});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["value"], serde_json::json!("{\"mode\":\"raw\"}")); // safety: test-only assertion
}
#[test]
fn permissive_schema_is_noop() {
let schema = serde_json::json!({
"type": "object",
"properties": {},
"additionalProperties": true
});
let params = serde_json::json!({"count": "10"});
let result = prepare_params_for_schema(&params, &schema);
assert_eq!(result["count"], serde_json::json!("10")); // safety: test-only assertion
}
#[test]
fn prepare_tool_params_uses_discovery_schema() {
let tool = StubTool {
schema: serde_json::json!({
"type": "object",
"properties": {
"requests": {
"type": "array",
"items": { "type": "object" }
}
}
}),
};
let params = serde_json::json!({
"requests": "[{\"insertText\":{\"text\":\"hello\"}}]"
});
let result = prepare_tool_params(&tool, &params);
#[rustfmt::skip]
assert_eq!( // safety: test-only assertion
result["requests"],
serde_json::json!([{ "insertText": { "text": "hello" } }])
);
}
}
+59 -4
View File
@@ -8,7 +8,7 @@ use crate::context::JobContext;
use crate::error::Error;
use crate::llm::ChatMessage;
use crate::safety::SafetyLayer;
use crate::tools::{ToolRegistry, redact_params};
use crate::tools::{ToolRegistry, prepare_tool_params, redact_params};
/// Execute a tool with safety checks: lookup → validate → timeout → execute → serialize.
///
@@ -29,8 +29,10 @@ pub async fn execute_tool_with_safety(
name: tool_name.to_string(),
})?;
let normalized_params = prepare_tool_params(tool.as_ref(), params);
// Validate tool parameters
let validation = safety.validator().validate_tool_params(params);
let validation = safety.validator().validate_tool_params(&normalized_params);
if !validation.is_valid {
let details = validation
.errors
@@ -45,7 +47,7 @@ pub async fn execute_tool_with_safety(
.into());
}
let safe_params = redact_params(params, tool.sensitive_params());
let safe_params = redact_params(&normalized_params, tool.sensitive_params());
tracing::debug!(
tool = %tool_name,
params = %safe_params,
@@ -56,7 +58,7 @@ pub async fn execute_tool_with_safety(
let timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(timeout, async {
tool.execute(params.clone(), job_ctx).await
tool.execute(normalized_params.clone(), job_ctx).await
})
.await;
let elapsed = start.elapsed();
@@ -237,6 +239,39 @@ mod tests {
}
}
struct ArrayEchoTool;
#[async_trait::async_trait]
impl Tool for ArrayEchoTool {
fn name(&self) -> &str {
"array_echo"
}
fn description(&self) -> &str {
"Echoes normalized params"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"values": {
"type": "array",
"items": { "type": "integer" }
}
}
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
Ok(ToolOutput::success(params, Duration::default()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
fn test_safety() -> SafetyLayer {
SafetyLayer::new(&crate::config::SafetyConfig {
max_output_length: 100_000,
@@ -348,6 +383,26 @@ mod tests {
);
}
#[tokio::test]
async fn test_execute_normalizes_stringified_array_params() {
let registry = registry_with(vec![Arc::new(ArrayEchoTool)]).await;
let safety = test_safety();
let result = execute_tool_with_safety(
&registry,
&safety,
"array_echo",
&serde_json::json!({"values": "[\"1\", \"2\", 3]"}),
&test_job_ctx(),
)
.await
.expect("array_echo should succeed"); // safety: test-only assertion
let output: serde_json::Value =
serde_json::from_str(&result).expect("tool result should be valid JSON"); // safety: test-only assertion
assert_eq!(output["values"], serde_json::json!([1, 2, 3])); // safety: test-only assertion
}
#[test]
fn test_process_tool_result_success() {
let safety = test_safety();
+2
View File
@@ -9,6 +9,7 @@
pub mod builder;
pub mod builtin;
mod coercion;
pub mod execute;
pub mod mcp;
pub mod rate_limiter;
@@ -24,6 +25,7 @@ pub use builder::{
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
};
pub(crate) use coercion::prepare_tool_params;
pub use rate_limiter::RateLimiter;
pub use registry::ToolRegistry;
pub use tool::{
+127 -164
View File
@@ -485,7 +485,7 @@ struct WasmToolSchemas {
/// 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.
/// Full schema available for discovery and runtime parameter preparation.
///
/// Seeded from the WASM `schema()` export at registration time, unless a
/// sidecar explicitly overrides it.
@@ -508,6 +508,19 @@ impl WasmToolSchemas {
.is_none_or(|p| p.is_empty())
}
fn typed_property_count(schema: &serde_json::Value) -> usize {
schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| {
props
.values()
.filter(|prop| schema_is_typed_property(prop))
.count()
})
.unwrap_or(0)
}
fn new(discovery: serde_json::Value) -> Self {
Self {
advertised: Self::permissive_schema(),
@@ -533,27 +546,6 @@ impl WasmToolSchemas {
fn discovery(&self) -> serde_json::Value {
self.discovery.clone()
}
/// Return the best schema available for type coercion.
///
/// Prefers the discovery schema when it has typed properties. Falls back
/// to the `PreparedModule` schema extracted at load time rather than
/// re-calling the WASM `schema()` export mid-execution, which could
/// interact with mutable linear memory state.
fn effective_for_coercion(&self, prepared_schema: &serde_json::Value) -> serde_json::Value {
if !Self::is_permissive_schema(&self.discovery) {
return self.discovery.clone();
}
// Fall back to the load-time extracted schema from PreparedModule.
// This avoids calling schema() on the already-running WASM instance
// where mutable state could produce inconsistent results.
if !Self::is_permissive_schema(prepared_schema) {
return prepared_schema.clone();
}
self.discovery.clone()
}
}
impl WasmToolWrapper {
@@ -583,7 +575,21 @@ impl WasmToolWrapper {
/// Override the parameter schema.
pub fn with_schema(mut self, schema: serde_json::Value) -> Self {
self.schemas = self.schemas.with_override(schema);
let override_typed = WasmToolSchemas::typed_property_count(&schema);
let prepared_typed = WasmToolSchemas::typed_property_count(&self.prepared.schema);
if override_typed == 0 && prepared_typed > 0 {
tracing::warn!(
tool = %self.prepared.name,
"Ignoring untyped schema override for discovery/runtime preparation and preserving extracted WASM schema"
);
self.schemas = WasmToolSchemas {
advertised: schema,
discovery: self.prepared.schema.clone(),
};
} else {
self.schemas = self.schemas.with_override(schema);
}
self
}
@@ -697,16 +703,6 @@ impl WasmToolWrapper {
// Get typed interface — used for execute.
let tool_iface = instance.near_agent_tool();
// Determine effective schema for type coercion.
// Prefer the discovery schema when typed; fall back to the load-time
// extracted schema from PreparedModule rather than re-calling the WASM
// export on the already-running instance.
let effective_schema = self.schemas.effective_for_coercion(&self.prepared.schema);
// 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, &effective_schema);
// Prepare the request
let params_json = serde_json::to_string(&params)
.map_err(|e| WasmError::InvalidResponseJson(e.to_string()))?;
@@ -734,10 +730,7 @@ impl WasmToolWrapper {
// 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 = format!(
"Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.",
self.prepared.name
);
let hint = build_tool_usage_hint(&self.prepared.name, &self.schemas.discovery());
return Err(WasmError::ToolReturnedError { message: err, hint });
}
@@ -1325,59 +1318,69 @@ fn is_private_ip(ip: std::net::IpAddr) -> bool {
}
}
/// Coerce parameter values to match their JSON Schema-declared types.
///
/// LLMs frequently send numeric values as strings (e.g. `"5"` instead of `5`)
/// or booleans as strings (`"true"` instead of `true`). This walks the params
/// object and converts string values where the schema expects a different type.
fn coerce_params_to_schema(
mut params: serde_json::Value,
schema: &serde_json::Value,
) -> serde_json::Value {
let properties = schema.get("properties").and_then(|p| p.as_object());
fn schema_contains_container_properties(schema: &serde_json::Value) -> bool {
schema
.get("properties")
.and_then(|p| p.as_object())
.map(|props| {
props.values().any(|prop| {
schema_declares_type(prop, "array") || schema_declares_type(prop, "object")
})
})
.unwrap_or(false)
}
let properties = match properties {
Some(p) => p,
None => return params,
};
let obj = match params.as_object_mut() {
Some(o) => o,
None => return params,
};
for (key, prop_schema) in properties {
let declared_type = prop_schema.get("type").and_then(|t| t.as_str());
let declared_type = match declared_type {
Some(t) => t,
None => continue,
};
if let Some(current_value) = obj.get_mut(key)
&& let Some(s) = current_value.as_str()
{
if declared_type == "string" {
continue;
fn schema_declares_type(schema: &serde_json::Value, expected: &str) -> bool {
match schema.get("type") {
Some(serde_json::Value::String(t)) => t == expected,
Some(serde_json::Value::Array(types)) => types.iter().any(|t| t.as_str() == Some(expected)),
_ => match expected {
"object" => {
schema
.get("properties")
.and_then(|p| p.as_object())
.is_some()
|| schema
.get("additionalProperties")
.is_some_and(serde_json::Value::is_object)
}
"array" => schema.get("items").is_some(),
_ => false,
},
}
}
let coerced = match declared_type {
"number" => s.parse::<f64>().ok().map(serde_json::Value::from),
"integer" => s.parse::<i64>().ok().map(serde_json::Value::from),
"boolean" => match s.to_lowercase().as_str() {
"true" => Some(serde_json::json!(true)),
"false" => Some(serde_json::json!(false)),
_ => None,
},
_ => None,
};
fn schema_is_typed_property(schema: &serde_json::Value) -> bool {
matches!(
schema.get("type"),
Some(serde_json::Value::String(_)) | Some(serde_json::Value::Array(_))
) || schema.get("$ref").is_some()
|| schema.get("anyOf").is_some()
|| schema.get("oneOf").is_some()
|| schema.get("allOf").is_some()
|| schema.get("items").is_some()
|| schema
.get("properties")
.and_then(|p| p.as_object())
.is_some()
|| schema
.get("additionalProperties")
.is_some_and(serde_json::Value::is_object)
}
if let Some(new_val) = coerced {
*current_value = new_val;
}
}
fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String {
let mut hint = format!(
"Tip: call tool_info(name: \"{}\", include_schema: true) for the full parameter schema.",
tool_name
);
if schema_contains_container_properties(schema) {
hint.push_str(
" For array/object fields, pass native JSON arrays/objects, not quoted JSON strings.",
);
}
params
hint
}
#[cfg(test)]
@@ -1945,100 +1948,60 @@ mod tests {
assert!(result.is_ok());
}
#[test]
fn test_coerce_params_string_to_number() {
let schema = serde_json::json!({
#[tokio::test]
async fn test_untyped_override_preserves_extracted_discovery_schema() {
let typed_schema = serde_json::json!({
"type": "object",
"properties": {
"count": { "type": "number" },
"name": { "type": "string" }
"values": {
"type": ["array", "null"],
"items": { "type": "array" }
}
}
});
let params = serde_json::json!({"count": "5", "name": "test"});
let result = super::coerce_params_to_schema(params, &schema);
assert_eq!(result["count"], serde_json::json!(5.0));
assert_eq!(result["name"], serde_json::json!("test"));
let runtime = Arc::new(WasmToolRuntime::new(WasmRuntimeConfig::for_testing()).unwrap()); // safety: test-only setup
let mut prepared = runtime
.prepare("sheets", b"\0asm\x0d\0\x01\0", None)
.await
.unwrap(); // safety: test-only setup
Arc::get_mut(&mut prepared).unwrap().schema = typed_schema.clone(); // safety: test-only setup
let wrapper =
super::WasmToolWrapper::new(Arc::clone(&runtime), prepared, Capabilities::default())
.with_schema(serde_json::json!({
"type": "object",
"properties": {},
"additionalProperties": true
}));
#[rustfmt::skip]
assert_eq!( // safety: test-only assertion
wrapper.parameters_schema(),
serde_json::json!({
"type": "object",
"properties": {},
"additionalProperties": true
})
);
assert_eq!(wrapper.discovery_schema(), typed_schema); // safety: test-only assertion
}
#[test]
fn test_coerce_params_string_to_integer() {
fn test_build_tool_usage_hint_detects_nullable_container_properties() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"limit": { "type": "integer" }
"requests": {
"type": ["array", "null"],
"items": { "type": "object" }
}
}
});
let params = serde_json::json!({"limit": "10"});
let result = super::coerce_params_to_schema(params, &schema);
assert_eq!(result["limit"], serde_json::json!(10));
}
#[test]
fn test_coerce_params_string_to_boolean() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"a": { "type": "boolean" },
"b": { "type": "boolean" },
"c": { "type": "boolean" },
"d": { "type": "boolean" }
}
});
let params = serde_json::json!({
"a": "true",
"b": "false",
"c": "True",
"d": "FALSE"
});
let result = super::coerce_params_to_schema(params, &schema);
assert_eq!(result["a"], serde_json::json!(true));
assert_eq!(result["b"], serde_json::json!(false));
assert_eq!(result["c"], serde_json::json!(true));
assert_eq!(result["d"], serde_json::json!(false));
}
let hint = super::build_tool_usage_hint("google_docs", &schema);
#[test]
fn test_coerce_params_already_correct_type() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"count": { "type": "number" }
}
});
let params = serde_json::json!({"count": 5});
let result = super::coerce_params_to_schema(params, &schema);
assert_eq!(result["count"], serde_json::json!(5));
}
#[test]
fn test_coerce_params_invalid_string_not_coerced() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"count": { "type": "number" }
}
});
let params = serde_json::json!({"count": "not-a-number"});
let result = super::coerce_params_to_schema(params, &schema);
// Should remain as string since it can't be parsed
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"));
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
}
/// Regression test: leak scan must run on raw headers (before credential
+27 -16
View File
@@ -30,7 +30,7 @@ use crate::llm::{
use crate::safety::SafetyLayer;
use crate::tools::execute::process_tool_result;
use crate::tools::rate_limiter::RateLimitResult;
use crate::tools::{ApprovalContext, ToolRegistry, redact_params};
use crate::tools::{ApprovalContext, ToolRegistry, prepare_tool_params, redact_params};
/// Shared dependencies for worker execution.
///
@@ -483,8 +483,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
name: tool_name.to_string(),
})?;
let normalized_params = prepare_tool_params(tool.as_ref(), params);
// Check approval: use context-aware check if available, else block all non-Never tools
let requirement = tool.requires_approval(params);
let requirement = tool.requires_approval(&normalized_params);
let blocked =
ApprovalContext::is_blocked_or_default(&deps.approval_context, tool_name, requirement);
if blocked {
@@ -517,9 +519,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
// Run BeforeToolCall hook
let params = {
let effective_params = {
use crate::hooks::{HookError, HookEvent, HookOutcome};
let hook_params = redact_params(params, tool.sensitive_params());
let hook_params = redact_params(&normalized_params, tool.sensitive_params());
let event = HookEvent::ToolCall {
tool_name: tool_name.to_string(),
parameters: hook_params,
@@ -543,15 +545,21 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
Ok(HookOutcome::Continue {
modified: Some(new_params),
}) => serde_json::from_str(&new_params).unwrap_or_else(|e| {
tracing::warn!(
tool = %tool_name,
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
e
);
params.clone()
}),
_ => params.clone(),
}) => match serde_json::from_str(&new_params) {
// Hook output is fresh JSON text and may reintroduce stringified scalars or
// containers, so we normalize it again. The fallback path reuses the already
// normalized input because no hook mutation was applied.
Ok(parsed) => prepare_tool_params(tool.as_ref(), &parsed),
Err(e) => {
tracing::warn!(
tool = %tool_name,
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
e
);
normalized_params
}
},
_ => normalized_params,
}
};
if job_ctx.state == JobState::Cancelled {
@@ -563,7 +571,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
// Validate tool parameters
let validation = deps.safety.validator().validate_tool_params(&params);
let validation = deps
.safety
.validator()
.validate_tool_params(&effective_params);
if !validation.is_valid {
let details = validation
.errors
@@ -579,7 +590,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
// Redact sensitive parameter values before they touch any observability or audit path.
let safe_params = redact_params(&params, tool.sensitive_params());
let safe_params = redact_params(&effective_params, tool.sensitive_params());
tracing::debug!(
tool = %tool_name,
params = %safe_params,
@@ -591,7 +602,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let tool_timeout = tool.execution_timeout();
let start = std::time::Instant::now();
let result = tokio::time::timeout(tool_timeout, async {
tool.execute(params.clone(), &job_ctx).await
tool.execute(effective_params.clone(), &job_ctx).await
})
.await;
let elapsed = start.elapsed();
+346
View File
@@ -0,0 +1,346 @@
//! E2E trace tests: schema-guided tool parameter normalization.
//!
//! These regressions run through the real agent loop with stub tools that
//! mirror Google Sheets / Google Docs write payload shapes. The model sends
//! quoted JSON container values, and the runtime must normalize them before
//! tool execution.
#[cfg(feature = "libsql")]
mod support;
#[cfg(feature = "libsql")]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use serde_json::json;
use ironclaw::context::JobContext;
use ironclaw::tools::{Tool, ToolError, ToolOutput};
use crate::support::test_rig::TestRigBuilder;
use crate::support::trace_llm::{
LlmTrace, TraceExpects, TraceResponse, TraceStep, TraceToolCall,
};
struct SheetsWriteFixtureTool;
#[async_trait]
impl Tool for SheetsWriteFixtureTool {
fn name(&self) -> &str {
"google_sheets_write_fixture"
}
fn description(&self) -> &str {
"Test fixture for Sheets-style values writes"
}
fn parameters_schema(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"spreadsheet_id": { "type": "string" },
"range": { "type": "string" },
"values": {
"type": "array",
"items": {
"type": "array",
"items": { "type": "integer" }
}
}
},
"required": ["spreadsheet_id", "range", "values"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let rows = params
.get("values")
.and_then(|v| v.as_array())
.ok_or_else(|| ToolError::InvalidParameters("values must be an array".into()))?;
let mut sum = 0_i64;
for row in rows {
let cells = row.as_array().ok_or_else(|| {
ToolError::InvalidParameters("each row must be an array".into())
})?;
for cell in cells {
sum += cell.as_i64().ok_or_else(|| {
ToolError::InvalidParameters("all cells must be integers".into())
})?;
}
}
Ok(ToolOutput::success(
json!({
"rows": rows.len(),
"sum": sum
}),
Duration::from_millis(1),
))
}
fn requires_sanitization(&self) -> bool {
false
}
}
struct DocsBatchUpdateFixtureTool;
#[async_trait]
impl Tool for DocsBatchUpdateFixtureTool {
fn name(&self) -> &str {
"google_docs_batch_update_fixture"
}
fn description(&self) -> &str {
"Test fixture for Docs-style batchUpdate requests"
}
fn parameters_schema(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"document_id": { "type": "string" },
"requests": {
"type": "array",
"items": {
"type": "object",
"properties": {
"insert_text": {
"type": "object",
"properties": {
"location": {
"type": "object",
"properties": {
"index": { "type": "integer" }
},
"required": ["index"]
},
"text": { "type": "string" },
"bold": { "type": "boolean" }
},
"required": ["location", "text", "bold"]
}
},
"required": ["insert_text"]
}
}
},
"required": ["document_id", "requests"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let requests = params
.get("requests")
.and_then(|v| v.as_array())
.ok_or_else(|| ToolError::InvalidParameters("requests must be an array".into()))?;
let mut indexes = Vec::new();
let mut bold_count = 0_usize;
for request in requests {
let insert = request
.get("insert_text")
.and_then(|v| v.as_object())
.ok_or_else(|| {
ToolError::InvalidParameters("insert_text must be an object".into())
})?;
let index = insert
.get("location")
.and_then(|v| v.get("index"))
.and_then(|v| v.as_i64())
.ok_or_else(|| {
ToolError::InvalidParameters("location.index must be an integer".into())
})?;
if insert
.get("bold")
.and_then(|v| v.as_bool())
.unwrap_or(false)
{
bold_count += 1;
}
indexes.push(index);
}
Ok(ToolOutput::success(
json!({
"request_count": requests.len(),
"indexes": indexes,
"bold_count": bold_count
}),
Duration::from_millis(1),
))
}
fn requires_sanitization(&self) -> bool {
false
}
}
#[tokio::test]
async fn e2e_normalizes_stringified_google_sheets_values() {
let trace = LlmTrace {
model_name: "test-coercion-sheets".to_string(),
turns: vec![crate::support::trace_llm::TraceTurn {
user_input: "Append these rows to the sheet".to_string(),
steps: vec![
TraceStep {
request_hint: None,
response: TraceResponse::ToolCalls {
tool_calls: vec![TraceToolCall {
id: "call_sheets".to_string(),
name: "google_sheets_write_fixture".to_string(),
arguments: json!({
"spreadsheet_id": "sheet-123",
"range": "Sheet1!A1:B2",
"values": "[[\"1\",2],[\"3\",\"4\"]]"
}),
}],
input_tokens: 100,
output_tokens: 25,
},
expected_tool_results: Vec::new(),
},
TraceStep {
request_hint: None,
response: TraceResponse::Text {
content: "The sheet write succeeded with 2 rows and sum 10."
.to_string(),
input_tokens: 120,
output_tokens: 20,
},
expected_tool_results: Vec::new(),
},
],
expects: TraceExpects::default(),
}],
memory_snapshot: Vec::new(),
http_exchanges: Vec::new(),
expects: TraceExpects {
response_contains: vec!["2 rows".to_string(), "sum 10".to_string()],
response_not_contains: Vec::new(),
response_matches: None,
tools_used: vec!["google_sheets_write_fixture".to_string()],
tools_not_used: Vec::new(),
all_tools_succeeded: Some(true),
max_tool_calls: Some(1),
min_responses: Some(1),
tool_results_contain: std::collections::HashMap::new(),
tools_order: vec!["google_sheets_write_fixture".to_string()],
},
steps: Vec::new(),
};
let rig = TestRigBuilder::new()
.with_trace(trace.clone())
.with_extra_tools(vec![Arc::new(SheetsWriteFixtureTool)])
.build()
.await;
rig.send_message("Append these rows to the sheet").await;
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
rig.verify_trace_expects(&trace, &responses);
let tool_results = rig.tool_results();
assert!(
tool_results
.iter()
.any(|(name, preview)| name == "google_sheets_write_fixture"
&& preview.contains("\"rows\"")
&& preview.contains("2")
&& preview.contains("\"sum\"")
&& preview.contains("10")),
"expected normalized sheet result preview, got {tool_results:?}"
);
rig.shutdown();
}
#[tokio::test]
async fn e2e_normalizes_stringified_google_docs_requests() {
let trace = LlmTrace {
model_name: "test-coercion-docs".to_string(),
turns: vec![crate::support::trace_llm::TraceTurn {
user_input: "Apply these edits to the doc".to_string(),
steps: vec![
TraceStep {
request_hint: None,
response: TraceResponse::ToolCalls {
tool_calls: vec![TraceToolCall {
id: "call_docs".to_string(),
name: "google_docs_batch_update_fixture".to_string(),
arguments: json!({
"document_id": "doc-456",
"requests": "[{\"insert_text\":{\"location\":{\"index\":\"1\"},\"text\":\"Hello\",\"bold\":\"true\"}},{\"insert_text\":{\"location\":{\"index\":5},\"text\":\" world\",\"bold\":\"false\"}}]"
}),
}],
input_tokens: 140,
output_tokens: 30,
},
expected_tool_results: Vec::new(),
},
TraceStep {
request_hint: None,
response: TraceResponse::Text {
content: "The doc update succeeded with 2 requests at indexes 1 and 5."
.to_string(),
input_tokens: 180,
output_tokens: 24,
},
expected_tool_results: Vec::new(),
},
],
expects: TraceExpects::default(),
}],
memory_snapshot: Vec::new(),
http_exchanges: Vec::new(),
expects: TraceExpects {
response_contains: vec!["2 requests".to_string(), "indexes 1 and 5".to_string()],
response_not_contains: Vec::new(),
response_matches: None,
tools_used: vec!["google_docs_batch_update_fixture".to_string()],
tools_not_used: Vec::new(),
all_tools_succeeded: Some(true),
max_tool_calls: Some(1),
min_responses: Some(1),
tool_results_contain: std::collections::HashMap::new(),
tools_order: vec!["google_docs_batch_update_fixture".to_string()],
},
steps: Vec::new(),
};
let rig = TestRigBuilder::new()
.with_trace(trace.clone())
.with_extra_tools(vec![Arc::new(DocsBatchUpdateFixtureTool)])
.build()
.await;
rig.send_message("Apply these edits to the doc").await;
let responses = rig.wait_for_responses(1, Duration::from_secs(15)).await;
rig.verify_trace_expects(&trace, &responses);
let tool_results = rig.tool_results();
assert!(
tool_results
.iter()
.any(|(name, preview)| name == "google_docs_batch_update_fixture"
&& preview.contains("\"request_count\"")
&& preview.contains("2")
&& preview.contains("\"bold_count\"")
&& preview.contains("1")),
"expected normalized docs result preview, got {tool_results:?}"
);
rig.shutdown();
}
}