From fda5160940b5661dc85150dc1bbedd3e1ca6b8fc Mon Sep 17 00:00:00 2001 From: Henry Park Date: Sat, 14 Mar 2026 16:26:39 -0700 Subject: [PATCH 1/2] Make no-panics CI check test-aware (#1160) * Make no-panics check test-aware * Handle proc-macro test attrs in no-panics check * Pin Python for no-panics CI job --- .github/workflows/code_style.yml | 47 +--- scripts/check_no_panics.py | 360 +++++++++++++++++++++++++++++++ 2 files changed, 364 insertions(+), 43 deletions(-) create mode 100644 scripts/check_no_panics.py diff --git a/.github/workflows/code_style.yml b/.github/workflows/code_style.yml index 705f261b..f89161d9 100644 --- a/.github/workflows/code_style.yml +++ b/.github/workflows/code_style.yml @@ -86,52 +86,13 @@ jobs: uses: actions/checkout@v6 with: fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" - name: Check for .unwrap(), .expect(), assert!() in production code run: | BASE="${{ github.event.pull_request.base.sha }}" - # Get the full diff for .rs files (production only, exclude tests/ directory) - DIFF=$(git diff "$BASE"...HEAD -- 'src/**/*.rs' 'crates/**/*.rs' || true) - - if [ -z "$DIFF" ]; then - echo "No production Rust changes detected." - exit 0 - fi - - # Extract added lines, skipping those inside test modules. - # Track whether we're inside a test module by watching hunk headers - # (lines starting with @@) whose context contains "mod tests" or "#[cfg(test)]". - ADDED=$(echo "$DIFF" | awk ' - /^@@/ { - # Hunk context (after the second @@) tells us the function/module scope - in_test = (tolower($0) ~ /mod tests/ || $0 ~ /#\[cfg\(test\)\]/ || $0 ~ /#\[test\]/) - } - /^\+[^+]/ && !in_test { print } - ' || true) - - if [ -z "$ADDED" ]; then - echo "No production Rust changes detected (test-only changes excluded)." - exit 0 - fi - - # Match panic-inducing patterns, excluding safety suppressions - VIOLATIONS=$(echo "$ADDED" \ - | grep -E '\.(unwrap|expect)\(|[^_]assert(_eq|_ne)?!' \ - | grep -Ev 'debug_assert|// safety:' \ - || true) - - if [ -n "$VIOLATIONS" ]; then - echo "::error::Found .unwrap(), .expect(), or assert!() in production code." - echo "Production code must use proper error handling instead of panicking." - echo "Suppress false positives with an inline '// safety: ' comment." - echo "" - echo "$VIOLATIONS" | head -20 - echo "" - COUNT=$(echo "$VIOLATIONS" | wc -l | tr -d ' ') - echo "Total: $COUNT violation(s)" - exit 1 - fi - - echo "OK: No panic-inducing calls in changed production code." + python3 scripts/check_no_panics.py --base "$BASE" --head HEAD # Roll-up job for branch protection code-style: diff --git a/scripts/check_no_panics.py b/scripts/check_no_panics.py new file mode 100644 index 00000000..55b90d21 --- /dev/null +++ b/scripts/check_no_panics.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +# Requires Python 3.10+ for PEP 604 union syntax such as `int | None`. + +import argparse +import pathlib +import re +import subprocess +import sys +import unittest +from dataclasses import dataclass + + +PANIC_PATTERN = re.compile(r"\.(?:unwrap|expect)\(|(? str: + result = subprocess.run( + ["git", *args], + check=True, + capture_output=True, + text=True, + ) + return result.stdout + + +def sanitize_line(line: str, state: LexerState) -> str: + chars = list(line) + out = [" "] * len(chars) + i = 0 + + while i < len(chars): + ch = chars[i] + nxt = chars[i + 1] if i + 1 < len(chars) else "" + + if state.block_comment_depth: + if ch == "/" and nxt == "*": + state.block_comment_depth += 1 + i += 2 + continue + if ch == "*" and nxt == "/": + state.block_comment_depth -= 1 + i += 2 + continue + i += 1 + continue + + if state.raw_string_hashes is not None: + if ch == '"': + hashes = 0 + j = i + 1 + while j < len(chars) and chars[j] == "#": + hashes += 1 + j += 1 + if hashes == state.raw_string_hashes: + state.raw_string_hashes = None + i = j + continue + i += 1 + continue + + if state.in_string: + if state.string_escape: + state.string_escape = False + elif ch == "\\": + state.string_escape = True + elif ch == '"': + state.in_string = False + i += 1 + continue + + if state.in_char: + if state.char_escape: + state.char_escape = False + elif ch == "\\": + state.char_escape = True + elif ch == "'": + state.in_char = False + i += 1 + continue + + if ch == "/" and nxt == "/": + break + if ch == "/" and nxt == "*": + state.block_comment_depth += 1 + i += 2 + continue + if ch == "r": + j = i + 1 + while j < len(chars) and chars[j] == "#": + j += 1 + if j < len(chars) and chars[j] == '"': + state.raw_string_hashes = j - i - 1 + i = j + 1 + continue + if ch == '"': + state.in_string = True + i += 1 + continue + if ch == "'": + # This can misclassify lifetimes like `'a` as char literals. That only + # risks false negatives by masking later code on the same line. + state.in_char = True + i += 1 + continue + + out[i] = ch + i += 1 + + return "".join(out) + + +def is_test_item(line: str, pending_test_attr: bool) -> tuple[bool, bool]: + match = ITEM_PATTERN.match(line) + if not match: + return False, False + + kind, name = match.groups() + named_tests_module = kind == "mod" and name == "tests" + return True, pending_test_attr or named_tests_module + + +def line_test_contexts(lines: list[str]) -> list[bool]: + contexts = [False] * len(lines) + lexer = LexerState() + block_stack: list[bool] = [] + pending_test_attr = False + pending_block_context: bool | None = None + + for idx, raw in enumerate(lines): + code = sanitize_line(raw, lexer) + stripped = code.strip() + current_context = block_stack[-1] if block_stack else False + + if TEST_ATTR_PATTERN.match(stripped): + pending_test_attr = True + + item_found, item_is_test = is_test_item(code, pending_test_attr) + if item_found: + pending_block_context = item_is_test or current_context + pending_test_attr = False + elif stripped and not stripped.startswith("#[") and pending_test_attr: + pending_test_attr = False + + contexts[idx] = current_context or bool(pending_block_context) + + for ch in code: + if ch == "{": + if pending_block_context is not None: + block_stack.append(pending_block_context) + pending_block_context = None + else: + block_stack.append(block_stack[-1] if block_stack else False) + elif ch == "}" and block_stack: + block_stack.pop() + + if stripped.endswith(";"): + pending_block_context = None + + return contexts + + +def changed_rust_files(base: str, head: str) -> list[pathlib.Path]: + output = run_git("diff", "--name-only", f"{base}...{head}", "--", "src", "crates") + files = [] + for line in output.splitlines(): + if line.endswith(".rs") and (line.startswith("src/") or line.startswith("crates/")): + files.append(pathlib.Path(line)) + return files + + +def added_lines_for_file(base: str, head: str, path: pathlib.Path) -> set[int]: + diff = run_git("diff", "--unified=0", f"{base}...{head}", "--", str(path)) + added: set[int] = set() + current_line = 0 + + for line in diff.splitlines(): + if line.startswith("@@"): + match = re.search(r"\+(\d+)(?:,(\d+))?", line) + if not match: + continue + current_line = int(match.group(1)) + continue + if line.startswith("+++ ") or line.startswith("--- "): + continue + if line.startswith("+"): + added.add(current_line) + current_line += 1 + elif line.startswith("-"): + continue + else: + current_line += 1 + + return added + + +def collect_violations(base: str, head: str) -> list[tuple[str, int, str]]: + violations: list[tuple[str, int, str]] = [] + + for path in changed_rust_files(base, head): + if not path.exists(): + continue + added_lines = added_lines_for_file(base, head, path) + if not added_lines: + continue + + lines = path.read_text(encoding="utf-8").splitlines() + contexts = line_test_contexts(lines) + lexer = LexerState() + sanitized = [sanitize_line(line, lexer) for line in lines] + + for line_no in sorted(added_lines): + if line_no < 1 or line_no > len(lines): + continue + if contexts[line_no - 1]: + continue + if "// safety:" in lines[line_no - 1]: + continue + if PANIC_PATTERN.search(sanitized[line_no - 1]): + violations.append((str(path), line_no, lines[line_no - 1].rstrip())) + + return violations + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--base", required=False, default="origin/staging") + parser.add_argument("--head", required=False, default="HEAD") + parser.add_argument("--self-test", action="store_true") + args = parser.parse_args() + + if args.self_test: + suite = unittest.defaultTestLoader.loadTestsFromTestCase(CheckNoPanicsTests) + result = unittest.TextTestRunner(verbosity=2).run(suite) + return 0 if result.wasSuccessful() else 1 + + violations = collect_violations(args.base, args.head) + if not violations: + print("OK: No panic-inducing calls in changed production code.") + return 0 + + print("::error::Found panic-style calls outside test-only Rust code.") + print("Production code must use proper error handling instead of panicking.") + print("Suppress false positives with an inline '// safety: ' comment.") + print("") + for path, line_no, line in violations[:20]: + print(f"{path}:{line_no}: {line}") + print("") + print(f"Total: {len(violations)} violation(s)") + return 1 + + +class CheckNoPanicsTests(unittest.TestCase): + def test_cfg_test_module_marks_inner_lines(self) -> None: + lines = [ + "#[cfg(test)]\n", + "mod tests {\n", + " assert!(true);\n", + "}\n", + "fn prod() {\n", + " value.expect(\"boom\");\n", + "}\n", + ] + + contexts = line_test_contexts(lines) + + self.assertTrue(contexts[1]) + self.assertTrue(contexts[2]) + self.assertFalse(contexts[4]) + self.assertFalse(contexts[5]) + + def test_test_function_marks_body_only(self) -> None: + lines = [ + "#[test]\n", + "fn it_works(\n", + ") {\n", + " assert_eq!(2 + 2, 4);\n", + "}\n", + "fn prod() {\n", + " assert!(ready);\n", + "}\n", + ] + + contexts = line_test_contexts(lines) + + self.assertTrue(contexts[1]) + self.assertTrue(contexts[2]) + self.assertTrue(contexts[3]) + self.assertFalse(contexts[5]) + self.assertFalse(contexts[6]) + + def test_proc_macro_test_attrs_mark_body_only(self) -> None: + attrs = [ + "tokio::test", + 'tokio::test(flavor = "multi_thread", worker_threads = 4)', + "rstest", + "test_case(1, 2)", + "cfg(all(test, unix))", + ] + + for attr in attrs: + with self.subTest(attr=attr): + lines = [ + f"#[{attr}]\n", + "fn it_works() {\n", + ' value.expect("allowed in test");\n', + "}\n", + "fn prod() {\n", + ' value.expect("boom");\n', + "}\n", + ] + + contexts = line_test_contexts(lines) + + self.assertTrue(contexts[1]) + self.assertTrue(contexts[2]) + self.assertFalse(contexts[4]) + self.assertFalse(contexts[5]) + + def test_named_tests_module_marks_context(self) -> None: + lines = [ + "mod tests {\n", + " fn helper() {\n", + " assert!(true);\n", + " }\n", + "}\n", + ] + + contexts = line_test_contexts(lines) + + self.assertTrue(all(contexts)) + + +if __name__ == "__main__": + sys.exit(main()) From c79754df2888ac7e2704d6cf4686b111eceee959 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Sat, 14 Mar 2026 16:27:18 -0700 Subject: [PATCH 2/2] 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 --- src/agent/routine_engine.rs | 14 +- src/agent/scheduler.rs | 87 +++++++- src/tools/builder/core.rs | 5 +- src/tools/coercion.rs | 367 +++++++++++++++++++++++++++++++ src/tools/execute.rs | 63 +++++- src/tools/mod.rs | 2 + src/tools/wasm/wrapper.rs | 291 +++++++++++------------- src/worker/job.rs | 43 ++-- tests/e2e_tool_param_coercion.rs | 346 +++++++++++++++++++++++++++++ 9 files changed, 1025 insertions(+), 193 deletions(-) create mode 100644 src/tools/coercion.rs create mode 100644 tests/e2e_tool_param_coercion.rs diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index 739b20d7..c37ba7ce 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -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(); diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 3923530f..fa7364a4 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -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 { + 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:?}" + ); + } } diff --git a/src/tools/builder/core.rs b/src/tools/builder/core.rs index 190fd21e..d4e10e95 100644 --- a/src/tools/builder/core.rs +++ b/src/tools/builder/core.rs @@ -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. diff --git a/src/tools/coercion.rs b/src/tools/coercion.rs new file mode 100644 index 00000000..34ef0057 --- /dev/null +++ b/src/tools/coercion.rs @@ -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 { + if schema_allows_type(schema, "string") { + return None; + } + + if schema_allows_type(schema, "integer") + && let Ok(v) = s.parse::() + { + return Some(serde_json::Value::from(v)); + } + + if schema_allows_type(schema, "number") + && let Ok(v) = s.parse::() + { + 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::(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 { + 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" } }]) + ); + } +} diff --git a/src/tools/execute.rs b/src/tools/execute.rs index 7c82d7ff..c6c20dc1 100644 --- a/src/tools/execute.rs +++ b/src/tools/execute.rs @@ -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 { + 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(); diff --git a/src/tools/mod.rs b/src/tools/mod.rs index e49cf396..d1659ddb 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -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::{ diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index d612cc46..a1b36548 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -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::().ok().map(serde_json::Value::from), - "integer" => s.parse::().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 diff --git a/src/worker/job.rs b/src/worker/job.rs index 86363f38..1247a552 100644 --- a/src/worker/job.rs +++ b/src/worker/job.rs @@ -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(); diff --git a/tests/e2e_tool_param_coercion.rs b/tests/e2e_tool_param_coercion.rs new file mode 100644 index 00000000..e5258762 --- /dev/null +++ b/tests/e2e_tool_param_coercion.rs @@ -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 { + 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 { + 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(); + } +}