mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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:
@@ -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
@@ -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(), ¶ms);
|
||||
|
||||
// Scheduler-specific approval check
|
||||
let requirement = tool.requires_approval(¶ms);
|
||||
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, ¶ms, &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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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(¶ms, &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(¶ms, &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(¶ms, &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(¶ms, &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(¶ms, &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(¶ms, &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(¶ms, &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(¶ms, &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, ¶ms);
|
||||
|
||||
#[rustfmt::skip]
|
||||
assert_eq!( // safety: test-only assertion
|
||||
result["requests"],
|
||||
serde_json::json!([{ "insertText": { "text": "hello" } }])
|
||||
);
|
||||
}
|
||||
}
|
||||
+59
-4
@@ -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(
|
||||
®istry,
|
||||
&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();
|
||||
|
||||
@@ -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
@@ -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(¶ms)
|
||||
.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
@@ -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(¶ms);
|
||||
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(¶ms, 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();
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user