From 1489cfbb5ebbedace8cf67cfb692a44b044f40fe Mon Sep 17 00:00:00 2001 From: cyq <15000851237@163.com> Date: Sun, 28 Jun 2026 14:41:15 +0800 Subject: [PATCH] fix(tui): reject malformed tool arguments --- crates/tui/src/core/engine.rs | 3 +- crates/tui/src/core/engine/dispatch.rs | 11 ++ crates/tui/src/core/engine/streaming.rs | 1 + crates/tui/src/core/engine/tests.rs | 13 ++- crates/tui/src/core/engine/turn_loop.rs | 14 +++ crates/tui/src/tools/arg_repair.rs | 24 ++-- .../tui/tests/features/tool_lifecycle.feature | 16 +++ crates/tui/tests/tool_lifecycle_acceptance.rs | 110 +++++++++++++++++- 8 files changed, 169 insertions(+), 23 deletions(-) diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 3273e7c51c..e0dd5049fb 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -3856,7 +3856,8 @@ use self::dispatch::should_parallelize_tool_batch; use self::dispatch::{ ParallelToolResult, ParallelToolResultEntry, ToolExecGuard, ToolExecOutcome, ToolExecutionBatch, ToolExecutionPlan, caller_allowed_for_tool, caller_type_for_tool_use, - final_tool_input, format_tool_error, mcp_tool_approval_description, mcp_tool_is_parallel_safe, + final_tool_input, format_tool_error, malformed_tool_arguments_error, + malformed_tool_arguments_input, mcp_tool_approval_description, mcp_tool_is_parallel_safe, mcp_tool_is_read_only, parse_parallel_tool_calls, parse_tool_input, plan_tool_execution_batches, should_force_update_plan_first, should_stop_after_plan_tool, }; diff --git a/crates/tui/src/core/engine/dispatch.rs b/crates/tui/src/core/engine/dispatch.rs index c883a6872d..4f85047c5d 100644 --- a/crates/tui/src/core/engine/dispatch.rs +++ b/crates/tui/src/core/engine/dispatch.rs @@ -273,6 +273,9 @@ fn is_transient_tool_failure(err: &ToolError, formatted_message: &str) -> bool { /// (the per-delta parser has already mirrored the most recent valid /// partial parse into `tool_state.input`). pub(super) fn final_tool_input(state: &ToolUseState) -> serde_json::Value { + if state.input_parse_error.is_some() { + return malformed_tool_arguments_input(&state.input_buffer); + } if !state.input_buffer.trim().is_empty() && let Some(parsed) = parse_tool_input(&state.input_buffer) { @@ -307,6 +310,14 @@ pub(super) fn parse_tool_input(buffer: &str) -> Option { .and_then(|segment| serde_json::from_str::(&segment).ok()) } +pub(super) fn malformed_tool_arguments_input(buffer: &str) -> serde_json::Value { + json!({ "raw_arguments": buffer }) +} + +pub(super) fn malformed_tool_arguments_error(buffer: &str) -> String { + format!("malformed tool arguments from model: expected valid JSON, got {buffer:?}") +} + fn strip_code_fences(text: &str) -> Option { if !text.contains("```") { return None; diff --git a/crates/tui/src/core/engine/streaming.rs b/crates/tui/src/core/engine/streaming.rs index 50c4b87b21..ab90a0ab6f 100644 --- a/crates/tui/src/core/engine/streaming.rs +++ b/crates/tui/src/core/engine/streaming.rs @@ -21,6 +21,7 @@ pub(super) struct ToolUseState { pub(super) input: serde_json::Value, pub(super) caller: Option, pub(super) input_buffer: String, + pub(super) input_parse_error: Option, } /// Maximum total bytes of text/thinking content before aborting the stream. diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index d76398c8c7..73dbce9b8e 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -5126,6 +5126,7 @@ fn tool_state(initial: serde_json::Value, buffer: &str) -> ToolUseState { input: initial, caller: None, input_buffer: buffer.into(), + input_parse_error: None, } } @@ -5147,11 +5148,13 @@ fn final_tool_input_falls_back_to_initial_when_buffer_empty() { } #[test] -fn final_tool_input_repairs_unparseable_buffer() { - // The arg_repair module converts unparseable input to an empty object - // {} so dispatch always proceeds. The buffer wins over the initial input. - let state = tool_state(json!({"command": "echo hi"}), "{not json"); - assert_eq!(final_tool_input(&state), json!({})); +fn final_tool_input_preserves_raw_buffer_for_parse_errors() { + let mut state = tool_state(json!({}), "{not json"); + state.input_parse_error = Some("malformed tool arguments".into()); + assert_eq!( + final_tool_input(&state), + json!({"raw_arguments": "{not json"}) + ); } // === #103 transparent stream-retry policy ===================================== diff --git a/crates/tui/src/core/engine/turn_loop.rs b/crates/tui/src/core/engine/turn_loop.rs index 64ff4751ac..225e4d193a 100644 --- a/crates/tui/src/core/engine/turn_loop.rs +++ b/crates/tui/src/core/engine/turn_loop.rs @@ -941,6 +941,7 @@ impl Engine { input, caller, input_buffer: String::new(), + input_parse_error: None, }); } ContentBlockStart::ServerToolUse { id, name, input } => { @@ -955,6 +956,7 @@ impl Engine { input, caller: None, input_buffer: String::new(), + input_parse_error: None, }); } }, @@ -1066,6 +1068,11 @@ impl Engine { "Tool '{}' failed to parse final input buffer: '{}'", tool_state.name, tool_state.input_buffer )); + let error = + malformed_tool_arguments_error(&tool_state.input_buffer); + tool_state.input_parse_error = Some(error); + tool_state.input = + malformed_tool_arguments_input(&tool_state.input_buffer); let _ = self .tx_event .send(Event::status(format!( @@ -1223,6 +1230,7 @@ impl Engine { input: call.args, caller: None, input_buffer: String::new(), + input_parse_error: None, }); } } @@ -1588,6 +1596,12 @@ impl Engine { ))); } + if blocked_error.is_none() + && let Some(error) = tool.input_parse_error.clone() + { + blocked_error = Some(ToolError::invalid_input(error)); + } + // #3027: deny wins over allow — check the deny-list first so a // tool present in both lists is still blocked. if blocked_error.is_none() diff --git a/crates/tui/src/tools/arg_repair.rs b/crates/tui/src/tools/arg_repair.rs index 170a78bfdc..2a7499f9de 100644 --- a/crates/tui/src/tools/arg_repair.rs +++ b/crates/tui/src/tools/arg_repair.rs @@ -5,16 +5,15 @@ //! reassembly leaves a trailing comma or unclosed brace; (b) some local //! backends emit literal control characters inside JSON string values. //! -//! The repair ladder runs five stages before falling back to an empty object: +//! The repair ladder runs five stages before reporting unrecoverable input: //! //! 1. Strict parse — done if it parses. //! 2. Strip literal control chars inside string values. //! 3. Strip trailing commas before `}` or `]`. //! 4. Balance braces/brackets (append closers). //! 5. Strip excess closers if delta is negative. -//! 6. Fallback: empty object `{}`. -use serde_json::{Map, Value}; +use serde_json::Value; /// Maximum raw argument length we'll attempt to repair (1 MiB). const MAX_ARG_LEN: usize = 1024 * 1024; @@ -23,12 +22,13 @@ const MAX_ARG_LEN: usize = 1024 * 1024; pub enum ArgRepairError { #[error("argument exceeded {0} chars; refusing to repair")] TooLarge(usize), + #[error("argument could not be repaired into valid JSON")] + Unrepairable, } /// Repair a raw JSON argument string into a valid `serde_json::Value`. /// /// Runs the deterministic ladder; on success returns the parsed value. -/// The final fallback is an empty object `{}` so dispatch always proceeds. pub fn repair(raw: &str) -> Result { if raw.len() > MAX_ARG_LEN { return Err(ArgRepairError::TooLarge(raw.len())); @@ -57,8 +57,7 @@ pub fn repair(raw: &str) -> Result { if let Ok(v) = serde_json::from_str(&s) { return Ok(v); } - // Fallback: empty object - Ok(Value::Object(Map::new())) + Err(ArgRepairError::Unrepairable) } /// Strip ASCII control characters (0x00–0x1F except \t, \n, \r) that appear @@ -224,15 +223,16 @@ mod tests { } #[test] - fn handles_empty_string() { - let v = repair("").unwrap(); - assert_eq!(v, json!({})); + fn rejects_empty_string() { + assert!(matches!(repair(""), Err(ArgRepairError::Unrepairable))); } #[test] - fn handles_gibberish() { - let v = repair("not json at all").unwrap(); - assert_eq!(v, json!({})); + fn rejects_gibberish() { + assert!(matches!( + repair("not json at all"), + Err(ArgRepairError::Unrepairable) + )); } #[test] diff --git a/crates/tui/tests/features/tool_lifecycle.feature b/crates/tui/tests/features/tool_lifecycle.feature index 2f250b001e..d1bcfa2149 100644 --- a/crates/tui/tests/features/tool_lifecycle.feature +++ b/crates/tui/tests/features/tool_lifecycle.feature @@ -51,3 +51,19 @@ Feature: Tool call lifecycle | status | marker | tool | input | | error | [!] | missing_tool | . | And the public output should include "I could not run the requested missing tool." + + Scenario: Malformed tool arguments return an error result + Given an offline CodeWhale workspace containing: + | path | kind | + | README.md | file | + And the mocked LLM will request the "list_dir" tool with malformed arguments "{not-json" + And the mocked LLM will answer after the tool result: + | content | + | I could not parse the tool arguments. | + When the user asks "try malformed tool arguments" + Then CodeWhale should send the user request to the mocked LLM + And the public tool lifecycle should show a running tool with raw input for "list_dir" + And the public tool result should report malformed arguments for "list_dir" + And CodeWhale should send the malformed argument error back to the mocked LLM + And the public tool lifecycle should show a failed tool with raw input for "list_dir" + And the public output should include "I could not parse the tool arguments." diff --git a/crates/tui/tests/tool_lifecycle_acceptance.rs b/crates/tui/tests/tool_lifecycle_acceptance.rs index e2fa9473c3..f2384a9cdd 100644 --- a/crates/tui/tests/tool_lifecycle_acceptance.rs +++ b/crates/tui/tests/tool_lifecycle_acceptance.rs @@ -19,6 +19,7 @@ const FEATURE_PATH: &str = concat!( ); const HAPPY_PATH_SCENARIO: &str = "Happy path lists the current directory through a tool"; const UNKNOWN_TOOL_SCENARIO: &str = "Unknown tool returns an error result"; +const MALFORMED_ARGUMENTS_SCENARIO: &str = "Malformed tool arguments return an error result"; const TOOL_CALL_ID: &str = "call_tool"; const TEST_MODEL: &str = "acceptance-model"; @@ -28,7 +29,7 @@ struct ToolLifecycleWorld { home: Option, llm_server: Option, tool_name: Option, - tool_input: Option, + tool_arguments: Option, final_answer: Option, prompt: Option, stdout: String, @@ -74,7 +75,19 @@ fn mocked_llm_will_request_tool(world: &mut ToolLifecycleWorld, tool_name: Strin ); world.tool_name = Some(tool_name); - world.tool_input = Some(input); + world.tool_arguments = Some(serde_json::to_string(&input).expect("tool input arguments")); +} + +#[given( + regex = r#"^the mocked LLM will request the "([^"]+)" tool with malformed arguments "([^"]+)"$"# +)] +fn mocked_llm_will_request_tool_with_malformed_arguments( + world: &mut ToolLifecycleWorld, + tool_name: String, + arguments: String, +) { + world.tool_name = Some(tool_name); + world.tool_arguments = Some(arguments); } #[given("the mocked LLM will answer after the tool result:")] @@ -243,6 +256,77 @@ fn codewhale_should_send_tool_error_back_to_mocked_llm(world: &mut ToolLifecycle ); } +#[then( + regex = r#"^the public tool lifecycle should show a running tool with raw input for "([^"]+)"$"# +)] +fn public_tool_lifecycle_should_show_running_tool_with_raw_input( + world: &mut ToolLifecycleWorld, + tool_name: String, +) { + let event = tool_use_event(world, &tool_name); + assert!( + value_contains_text(event.get("input").expect("tool_use input"), "{not-json"), + "tool_use input should preserve malformed raw arguments:\n{event:#}" + ); +} + +#[then(regex = r#"^the public tool result should report malformed arguments for "([^"]+)"$"#)] +fn public_tool_result_should_report_malformed_arguments_for( + world: &mut ToolLifecycleWorld, + tool_name: String, +) { + let _ = tool_use_event(world, &tool_name); + let event = tool_result_event(world); + + assert_eq!(event.get("status").and_then(Value::as_str), Some("error")); + let output = event + .get("output") + .and_then(Value::as_str) + .expect("tool_result error output"); + assert_malformed_arguments_text(output); +} + +#[then("CodeWhale should send the malformed argument error back to the mocked LLM")] +fn codewhale_should_send_malformed_argument_error_back_to_mocked_llm( + world: &mut ToolLifecycleWorld, +) { + let request = world + .requests + .iter() + .find(|request| request_contains_tool_result(request)) + .expect("expected a follow-up chat request containing the malformed argument error"); + let tool_result = tool_result_message(request).expect("tool result message"); + assert_eq!( + tool_result + .get("tool_call_id") + .and_then(serde_json::Value::as_str), + Some(TOOL_CALL_ID) + ); + + let content = tool_result + .get("content") + .and_then(serde_json::Value::as_str) + .expect("tool result content"); + assert_malformed_arguments_text(content); +} + +#[then( + regex = r#"^the public tool lifecycle should show a failed tool with raw input for "([^"]+)"$"# +)] +fn public_tool_lifecycle_should_show_failed_tool_with_raw_input( + world: &mut ToolLifecycleWorld, + tool_name: String, +) { + let event = tool_result_event(world); + assert_eq!(event.get("status").and_then(Value::as_str), Some("error")); + + let tool_use = tool_use_event(world, &tool_name); + assert!( + value_contains_text(tool_use.get("input").expect("tool_use input"), "{not-json"), + "failed tool_use input should preserve malformed raw arguments:\n{tool_use:#}" + ); +} + #[then("the public tool lifecycle should show a completed tool:")] fn public_tool_lifecycle_should_show_completed_tool(world: &mut ToolLifecycleWorld, step: &Step) { let expected = one_table_row(step); @@ -301,6 +385,11 @@ async fn unknown_tool_returns_error_result() { run_scenario(UNKNOWN_TOOL_SCENARIO, 10).await; } +#[tokio::test(flavor = "current_thread")] +async fn malformed_tool_arguments_return_error_result() { + run_scenario(MALFORMED_ARGUMENTS_SCENARIO, 10).await; +} + async fn run_scenario(name: &'static str, expected_steps: usize) { let writer = ToolLifecycleWorld::cucumber() .fail_on_skipped() @@ -344,7 +433,7 @@ async fn start_mock_llm(world: &ToolLifecycleWorld) -> MockServer { .and(request_has_no_tool_result) .respond_with(sse_response(&tool_call_sse( world.tool_name.as_ref().expect("tool name"), - world.tool_input.as_ref().expect("tool input"), + world.tool_arguments.as_ref().expect("tool arguments"), ))) .mount(&server) .await; @@ -479,8 +568,7 @@ fn preserve_host_env(command: &mut Command) { } } -fn tool_call_sse(tool_name: &str, tool_input: &Value) -> String { - let arguments = serde_json::to_string(tool_input).expect("tool input arguments"); +fn tool_call_sse(tool_name: &str, arguments: &str) -> String { [ sse_chunk(json!({ "id": "chatcmpl-tool", @@ -554,6 +642,18 @@ fn final_answer_sse(answer: &str) -> String { .join("") } +fn assert_malformed_arguments_text(text: &str) { + let lower = text.to_ascii_lowercase(); + assert!( + lower.contains("argument") + && (lower.contains("malformed") + || lower.contains("parse") + || lower.contains("json") + || lower.contains("invalid")), + "expected malformed argument error text:\n{text}" + ); +} + fn sse_chunk(value: Value) -> String { format!( "data: {}\n\n",